%% run_equilibrium_aggregates.m
% Aggregates-only pipeline. Writes CSVs with Z_f and occupation-level W_of, L_of.
% Reads:
%   "../../Data/ModelData/occupation_exposed_tasks.csv"
%       columns: soc, taskid, taskimportance, exposure
%   "../../Data/ModelData/occupation_distribution.csv"
%       columns: soc, naics, employmentshare, wagebillshare, wage

clear; clc;
do_full_ge = 1 ; %Keep equal to 1. Modifications of this parameter are no longer supported
use_parallel = true;
targetN = 7;  % how many workers do you want to use when parallelizing 

for type = 1:4  

elastalphaF = 0.0469 ; % Elasticity from Argente et al 
if type == 1
readcog = 1 ; % 1 = use cognitive task exposures. 0 = use eloundou task exposures 
scenario = "baseline"; 
end

if type == 2
readcog = 1 ; % 1 = use cognitive task exposures. 0 = use eloundou task exposures 
scenario = "baseline_no_product"; 
elastalphaF = 0; 
end

if type == 3
readcog = 0 ; % 1 = use cognitive task exposures. 0 = use eloundou task exposures 
scenario = "baseline"; % Baseline scenario, except Eloundou et al exposures 
end

if type == 4
readcog = 1 ; % 1 = use cognitive task exposures. 0 = use eloundou task exposures 
scenario = "baseline_with_spill"; 
end


%% ----------------- Parameters (defaults: your requested set) -----------------
params = struct();
params.beta  = 0.5; % <-- Baseline for beta is 0.5. But also fix at 0.25 or 0.75. 
if type == 6 
    params.beta = 0.75; 
end
params.zeta  = 0.42; % Fixed at 0.42 (BHM)  
params.nu    = 4.63;     % nu in (0, +inf), not equal to 1 <-- we will target nu 
target_sl = 0.6 ; %Karabarbounis and Neiman (2014) Figure 2
gammaj = (target_sl)^(1/params.nu) / ( (target_sl)^(1/params.nu) + (1-target_sl)^(1/params.nu));
params.a     = gammaj^params.nu;  % a_j = gamma_j^nu (target approx task labor share of 0.81)
params.b     = (1-gammaj)^params.nu;  % b_j = (1-gamma_j)^nu (target approx task capital share of 0.19)
params.factor_cost = (params.a+params.b)^(1/(params.nu-1));
params.chi   = 1.34; % Fixed at 1.34 (Caunedo)
params.psi   = 0.5; % Fixed at 0.5 (Humlum) 
params.theta = 1.34; 
if type == 4
    params.theta = 1.72 ; % Turns on spillovers 
end
params.M     = 1.0; % wage markdown, set to 1 with perfectly competitive markets

% GE labor supply parameters (used only for GE step)
params.phi   = 0.25*0;
params.sigma = 1;

% Log capital price change for exposed tasks (price falls)
% This is for 10-year horizon 
% Will also do for 20-year 
%conservative scenario 
if contains(scenario,"baseline")
delta = 10 * 0.1396; % <--- Might be good to allow for code to take in assumed price change. '
end

% Baseline file is cognitive_3cat 
% we will also do robustness to the file ending in _tasks_3cat.csv 
if readcog == 1
file_tasks = "../../Data/ModelData/occupation_exposed_tasks_cognitive_3cat.csv";
end
if readcog == 0 
file_tasks = "../../Data/ModelData/occupation_exposed_tasks_3cat.csv";
end
file_dist  = "../../Data/ModelData/occupation_distribution_consolidated.csv";
file_prod = "../../Data/ModelData/AI_product_patent_ratio_by_naics.csv"; 

% Use initial q as wage guess (benign)
use_q_as_initial_guess = true;

% Output dir
out_dir = "./outputs";
if ~exist(out_dir,'dir'); mkdir(out_dir); end

%% ----------------- Calibrate to get the appropriate gamma_j -----------------

%% ----------------- Load data -----------------
opts = detectImportOptions(file_dist, 'TextType', 'string'); % keep text as string
opts = setvartype(opts, 'naics', 'string');                  % force NAICS as text

Ttasks = readtable(file_tasks);   % soc, taskid, taskimportance, exposure
Tdist  = readtable(file_dist,opts);    % soc, naics, employmentshare, wagebillshare, wage
Tprod = readtable(file_prod,opts); % naics and mean_dlogai are the columns of interest

% Ensure string types (robust)
if ~isstring(Ttasks.soc);  Ttasks.soc  = string(Ttasks.soc);  end
if ~isstring(Tdist.soc);   Tdist.soc   = string(Tdist.soc);   end
if ~isstring(Tdist.naics); Tdist.naics = string(Tdist.naics); end

%% ----------------- Helpers -----------------
Wof_fun = @(w_vec, alphaJ, beta) ( sum( alphaJ(:) .* (w_vec(:).^(1/beta)) ) )^beta;

% %%% GE NEW: Lof uses Ybar (your GE mapping)
Lof_fun = @(Wof, zbar, Ybar, alphaOf, alphaF, zeta) alphaOf*alphaF*(Ybar .* zbar) .* (Wof.^zeta);

%% ----------------- Firms loop (BASELINE ONLY; store, don't write yet) -----------------
firms = unique(Tdist.naics);
numF  = numel(firms);

FirmRows = table('Size',[0 7], ...
    'VariableTypes', {'string','double','double','double','double','double', 'double'}, ...
    'VariableNames', {'naics','Z_f_before','Z_f_after','dlog_Z_f','dlog_Y_f','avg_dlog_q', 'dlog_alphaF'});

OccRows  = table('Size',[0 9], ...
    'VariableTypes', {'string','string','double','double','double','double','double','double','double'}, ...
    'VariableNames', {'naics','soc','W_of_before','W_of_after','dlog_W_of','L_of_before','L_of_after','dlog_L_of','dlog_L_of_fixedZZ'});

% --- Pick the first NAICS that actually has task rows ---
alphaJ = {}; alphaO = [];  found = false;
for f = 1:numel(firms)
    na = firms(f);
    Tdist_f = Tdist(Tdist.naics==na, :);
    socs_f  = unique(Tdist_f.soc);
    tmp_alphaJ = {}; tmp_alphaO = [];
    for o = 1:numel(socs_f)
        s = socs_f(o);
        Ts = Ttasks(Ttasks.soc==s, :);
        if isempty(Ts), continue; end           % skip SOCs with no tasks
        tmp_alphaJ{end+1,1} = Ts.taskimportance(:);
        tmp_alphaO(end+1,1) = Tdist_f.wagebillshare(Tdist_f.soc==s); %#ok<AGROW>
    end
    if ~isempty(tmp_alphaJ)
        alphaJ = tmp_alphaJ;
        alphaO = tmp_alphaO;
        target_naics = na;                       % for your logs
        found = true;
        break
    end
end
if ~found
    error('No NAICS found with at least one SOC that has tasks.');
end
alphaO = alphaO./sum(alphaO); % normalize weights

% --- Calibrate zbar once (uses α_j normalization internally) ---
cal = calibrate_zbar(alphaJ, alphaO, params.factor_cost, params.factor_cost, params); %, 'NormalizeAlphaO', false);
zbar_global = cal.zbar;
fprintf('Calibrated zbar = %.6f using NAICS %s\n', zbar_global, target_naics);

% --- Use this for everything that follows ---
params.zbar = zbar_global;

% %%% GE NEW: caches we’ll need for aggregates and the GE step
alphaOf_all         = cell(numF,1);   % α_{o,f}
alphaJ_all          = cell(numF,1);   % α_j lists per occupation
kept_socs_all       = cell(numF,1);
Q_init_all          = cell(numF,1);
Q_after_all         = cell(numF,1);
W_before_all        = cell(numF,1);   % cell of wages per occupation (for warm starts)
Wof_before_all      = cell(numF,1);   % vector of W_{o,f}
Zf_before_all       = nan(numF,1);
details_before_all  = cell(numF,1);   % needs patched solve_task_wages to expose Pof

% α_f (firm weights): use wagebillshare summed within NAICS (normalize)
alphaF = zeros(numF,1);
for f = 1:numF
    idx = Tdist.naics==firms(f);
    if any(ismember('indwagebillshare', Tdist.Properties.VariableNames))
        alphaF(f) = nansum(Tdist.indwagebillshare(idx));
    elseif any(ismember('indempshare', Tdist.Properties.VariableNames))
        alphaF(f) = nansum(Tdist.indempshare(idx));
    else
        alphaF(f) = nnz(idx);
    end
end
alphaF = alphaF / sum(alphaF);



for f = 1:numF
    na = firms(f);

    % SOCs present for this firm (allow different set per firm)
    Tdist_f = Tdist(Tdist.naics==na, :);
    socs_f  = unique(Tdist_f.soc);
    if isempty(socs_f), continue; end

    % Build α_o aligned to socs_f
    alphaO_full = zeros(numel(socs_f),1);
    for o = 1:numel(socs_f)
        row = Tdist_f(Tdist_f.soc==socs_f(o), :);
        alphaO_full(o) = row.wagebillshare(1);
    end

    % Build Q_init, alphaJ, exposure; SKIP SOCs with no tasks
    Q_init   = {};
    alphaJ   = {};
    exposure = {};
    kept_socs = strings(0,1);

    for o = 1:numel(socs_f)
        s = socs_f(o);
        Ttasks_s = Ttasks(Ttasks.soc==s, :);
        if isempty(Ttasks_s), continue; end     % Skip SOC appearing in dist but absent in tasks

        alphaJ{end+1,1}   = Ttasks_s.taskimportance(:);
        exposure{end+1,1} = Ttasks_s.exposure(:);
        Q_init{end+1,1}   = params.factor_cost.*ones(numel(Ttasks_s.taskimportance),1); % MATTHEW CHANGE
        kept_socs(end+1,1)= s;
    end

    % If nothing left after skipping, move on
    O = numel(Q_init);
    if O==0, continue; end

    % Shrink α_o to kept SOCs order and normalize
    alphaO = zeros(O,1);
    for o = 1:O
        row = Tdist_f(Tdist_f.soc==kept_socs(o), :);
        alphaO(o) = row.wagebillshare(1);
    end
    alphaO = alphaO./sum(alphaO);

    % Initial guess
    w0 = [];
    if use_q_as_initial_guess, w0 = Q_init; end

    % ====== BEFORE ======
    % (patched solve_task_wages returns details as 4th output)
    [W_before, Zf_before, ~, details_before] = solve_task_wages(Q_init, alphaJ, alphaO, alphaF(f), params, w0);

    % Cache baseline objects (no writes yet)
    W_before_all{f}       = W_before;
    alphaOf_all{f}        = alphaO;
    alphaJ_all{f}         = alphaJ;
    kept_socs_all{f}      = kept_socs;
    Q_init_all{f}         = Q_init;
    Zf_before_all(f)      = Zf_before;
    details_before_all{f} = details_before;

    % Build AFTER q shocked, but do not solve yet (GE step will handle it)
    Q_after = cell(O,1);
    for o = 1:O
        Q_after{o} = exp( -exposure{o} * delta ) .* Q_init{o};
    end
    Q_after_all{f} = Q_after;

    % Pre-compute W_{o,f}^{before}
    Wof_b_vec = zeros(O,1);
    for o = 1:O
        Wof_b_vec(o) = Wof_fun(W_before{o}, alphaJ{o}, params.beta);
    end
    Wof_before_all{f} = Wof_b_vec;
end
disp("Solved Initial wages")
%% ----------------- %%% GE NEW: Baseline aggregates and Φ -----------------

zeta  = params.zeta; chi = params.chi; theta = params.theta;
% \bar{W}_0
S_W = 0.0;
for f = 1:numF
    if isempty(alphaOf_all{f}), continue; end
    aF  = alphaF(f);
    aOf = alphaOf_all{f};
    Wof = Wof_before_all{f};
    for o = 1:numel(Wof)
        S_W = S_W + aF * aOf(o) * ( Wof(o)^(1+zeta) );
    end
end
Wbar0 = S_W^(1/(1+zeta));

% \bar{Y}_0 using P_{o,f} from details (fallback to Zf if needed)
% === Baseline \bar{Y}_0 via primal Y_f ===
theta = params.theta;
Yf0 = nan(numF,1);
for f = 1:numF
    if isempty(details_before_all{f}) || ~isfield(details_before_all{f},'Yf') || isempty(details_before_all{f}.Yf)
        % Fallback: compute directly if details missing
        [Yf0(f), ~] = Yf_from_wq(W_before_all{f}, Q_init_all{f}, alphaJ_all{f}, alphaOf_all{f}, alphaF(f),params);
    else
        Yf0(f) = details_before_all{f}.Yf;
    end
end
sF0   = sum( (alphaF(:).^(1/theta)) .* (Yf0(:).^((theta-1)/theta)) );
Ybar0 = sF0^(theta/(theta-1));


% \bar{N}_0 using new Lof_fun with Ybar0 and zbar0
zbar0 = params.zbar;

% Calibrate Φ from baseline identity:
% zbar = Wbar^(1+phi-zeta) * Φ * Nbar^(-phi*sigma) / Ybar
% where Nbar = Φ^( 1 / (1+sigma*phi) )  * Wbar^( phi*(1-sigma)) / (1+sigma*phi) ) 
phi   = params.phi;  sigma = params.sigma;
Phi = ( zbar0 * Ybar0 / ( (1 + sigma * phi)^(sigma * phi) ...
        * Wbar0^(1 + phi - zeta - phi^2 * sigma * (1 - sigma)) ) )^(1 + sigma * phi);
params.Phi = Phi ; 
fprintf('Calibrated Phi = %.8g\n', Phi);

%Equilibrium value of Nbar0 
Nbar0 =  Phi^(1/(1+sigma*phi)) * (Wbar0)^(  phi*(1-sigma) / (1+sigma*phi )    ) ; 

%-----------------------------
% Get average cost improvement 
%-----------------------------
avgDeltaQ = zeros(numel(alphaOf_all),1) ; 
for fff=1:numel(alphaOf_all) %Industries 
    alphaO_fff = alphaOf_all{fff}; 
    Q_after_fff = Q_after_all{fff}; 
    Q_before_fff = Q_init_all{fff}; 
    for iii=1:numel(alphaO_fff) %Occupations 
        alphaJiii = alphaJ_all{fff}{iii}; 
        for jjj=1:numel(alphaJiii) %Tasks 
            dlogQjjj = log(Q_after_fff{iii}(jjj)) - log(Q_before_fff{iii}(jjj)); 
            avgDeltaQ(fff) = avgDeltaQ(fff) + dlogQjjj * alphaO_fff(iii) * alphaJiii(jjj) ; 
        end
    end
end

%--------------------------------
% Allow for alphaF to respond to 
% industry AI innovation
%--------------------------------
alphaF_before = alphaF; 
dlogalphaF = Tprod.AI_ratio ; 
for fff=1:numel(dlogalphaF)
    alphaF(fff) = exp( elastalphaF * dlogalphaF(fff) ) * alphaF(fff) ;  
end 
%alphaF = alphaF / sum(alphaF);

%% ----------------- GE AFTER equilibrium (no local functions in script) -----------------
% Unknowns: gW = Wbar_after/Wbar0, gY = Ybar_after/Ybar0. gN is implied:
% gN = gW^(params.phi*(1-params.sigma))/(1+params.phi*params.sigma)).


%Starting a parallel pool 
if use_parallel && license('test','Distrib_Computing_Toolbox')
    p = gcp('nocreate');  % don't auto-create; just check
    try
        if isempty(p)
            % No pool yet: open one
            c = parcluster('local');
            parpool(c, min(targetN, c.NumWorkers));
        else
            % Pool exists: reuse if size matches, otherwise reopen
            if p.NumWorkers == targetN
                fprintf('Using existing parallel pool with %d workers.\n', p.NumWorkers);
            else
                fprintf('Reopening pool with %d workers (currently %d).\n', targetN, p.NumWorkers);
                delete(p);
                c = parcluster('local');                  % or: parcluster(p.Cluster.Profile)
                parpool(c, min(targetN, c.NumWorkers));
            end
        end
    catch ME
        warning('Parallel pool setup failed: %s\nFalling back to serial.');
        % (leave use_parallel=true if you want to retry later)
    end
end

%% Newton in (log gW, log gY) — drop-in replacement
x = [0; 0];              % [log gW; log gY]
tol = 1e-8;
maxit = 20;
h1 = 1e-4; h2 = 1e-4;    % FD steps

 

resid = @(xv) ge_resid_aftershock( ...
    xv, params, alphaF, alphaOf_all, alphaJ_all, ...
    Q_after_all, W_before_all, zbar0, Ybar0, Wbar0, Ybar0, Wof_fun);

if do_full_ge == 1
for it = 1:maxit
    r  = resid(x);
    if max(abs(r)) < tol, break; end

    % 2x2 finite-difference Jacobian
    e1 = [1;0]; e2 = [0;1];
    r1 = resid(x + h1*e1);
    r2 = resid(x + h2*e2);
    J  = [ (r1 - r)/h1, (r2 - r)/h2 ];

    % Newton step + line search
    dx  = -J\r;
    lam = 1.0; base = max(abs(r));
    for ls = 1:8
        xt = x + lam*dx;
        if max(abs(resid(xt))) < base, x = xt; break; end
        lam = lam/2;
    end
    if lam < 1/256, x = x + 1e-3*dx; end
end
end

[Wbar_after, Ybar_after, Nbar_after, Wof_after_all, Zf_after_all, Yf_after_all] = ge_forward_eval_aftershock( ...
    x, params, alphaF, alphaOf_all, alphaJ_all, Q_after_all, W_before_all, zbar0, Ybar0, Wof_fun);

gW = exp(x(1));
gY = exp(x(2));
gN = gW^(params.phi*(1-params.sigma)/(1+params.phi*params.sigma));

S  = gW^(1 + params.phi - params.zeta) * gN^(-params.phi*params.sigma) / gY;


%% ----------------- Build rows (now GE-consistent) -----------------
% Firm-level Z_f
for f = 1:numF
    if isnan(Zf_before_all(f)) || isnan(Zf_after_all(f)), continue; end
    FirmRows = [FirmRows; { ...
        firms(f), Zf_before_all(f), Zf_after_all(f), ....
        log(Zf_after_all(f)) - log(Zf_before_all(f)), ...
        log(Yf_after_all(f)) - log(Yf0(f)), avgDeltaQ(f),...
        elastalphaF*dlogalphaF(f)}]; %#ok<AGROW>
end

% === After-shock WITHOUT industry/economy GE effects: hold zbar and Z_f fixed at baseline ===
Wof_after_fixed_all = cell(numF,1);

for f = 1:numF
    if isempty(alphaOf_all{f}), continue; end

    params_fixed            = params;
    params_fixed.zbar       = zbar0;            % freeze zbar at baseline
    params_fixed.rhs_scale  = 1.0;              % no gY blue factor
    params_fixed.use_fixed_Zf = true;           % NEW flag
    params_fixed.Zf_fixed     = Zf_before_all(f);

    w0_after = W_before_all{f};                 % warm start from baseline

    [W_after_fixed, ~] = solve_task_wages( ...
        Q_after_all{f}, alphaJ_all{f}, alphaOf_all{f}, alphaF(f), params_fixed, w0_after);

    % Build W_{o,f}^{after, fixed}
    Of = numel(W_after_fixed);
    Wof_afix = zeros(Of,1);
    for o = 1:Of
        Wof_afix(o) = Wof_fun(W_after_fixed{o}, alphaJ_all{f}{o}, params.beta);
    end
    Wof_after_fixed_all{f} = Wof_afix;
end


% Occupation-level rows
% === After-shock WITHOUT industry/economy GE effects: hold zbar and Z_f fixed at baseline ===
disp("Solving with fixed Z_f and aggregates")
for f = 1:numF
    if isempty(alphaOf_all{f}), continue; end
    kept_socs = kept_socs_all{f};
    Wof_b_vec = Wof_before_all{f};
    Wof_a_vec = Wof_after_all{f};        % full GE
    Wof_fix   = Wof_after_fixed_all{f};  % fixed zbar & Z_f
    aF = alphaF(f); aF_before = alphaF_before(f);  
    alphaOf_vec = alphaOf_all{f}; 

    for o = 1:numel(Wof_b_vec)
        Wof_b = Wof_b_vec(o);
        Wof_a = Wof_a_vec(o);
        aOf = alphaOf_vec(o);  

        % Full-GE employment (already what you had)
        Lof_b      = Lof_fun(Wof_b, zbar0,       Ybar0,   aOf, aF_before,    params.zeta);
        Lof_a_full = Lof_fun(Wof_a, zbar0 * S,   Ybar0 * gY, aOf, aF, params.zeta);

        % NEW: “fixed aggregates” employment
        Lof_a_fix  = Lof_fun(Wof_fix(o), zbar0,  Ybar0, aOf, aF,    params.zeta);

        OccRows = [OccRows; { ...
            firms(f), kept_socs(o), ...
            Wof_b, Wof_a, log(Wof_a) - log(Wof_b), ...
            Lof_b, Lof_a_full, log(Lof_a_full) - log(Lof_b), ...
            log(Lof_a_fix) - log(Lof_b)}]; %#ok<AGROW>
    end
end

%% ----------------- Write CSVs -----------------
if readcog == 1
firm_csv = fullfile(out_dir, 'firm_Z_summary_' + scenario + '.csv');
occ_csv  = fullfile(out_dir, 'occupation_aggregates_' + scenario + '.csv');
end 
if readcog == 0
firm_csv = fullfile(out_dir, 'firm_Z_summary_eloundou_' + scenario + '.csv');
occ_csv  = fullfile(out_dir, 'occupation_aggregates_eloundou_' + scenario + '.csv');
end 

writetable(FirmRows, firm_csv);
writetable(OccRows,  occ_csv);

fprintf('Wrote %s\n', firm_csv);
fprintf('Wrote %s\n', occ_csv);

%% ----------------- Save aggregate (Wbar, Ybar, Nbar) summaries -----------------
AggRows = table( ...
    ["Wbar"; "Ybar"; "Nbar"], ...
    [Wbar0;  Ybar0;  Nbar0], ...
    [Wbar_after; Ybar_after; Nbar_after], ...
    [log(Wbar_after)-log(Wbar0);  log(Ybar_after)-log(Ybar0);  log(Nbar_after)-log(Nbar0)], ...
    [gW; gY; gN], ...
    'VariableNames', {'aggregate','baseline','after','dlog','growth'});

% Pick filename consistent with readcog
if readcog == 1
    agg_csv = fullfile(out_dir, 'aggregates_summary_' + scenario +'.csv');
else
    agg_csv = fullfile(out_dir, 'aggregates_summary_eloundou_' + scenario + '.csv');
end

writetable(AggRows, agg_csv);

fprintf('Wrote %s\n', agg_csv);

%% ----------------- Regression block (unchanged; now uses GE-consistent dlogs) -----------------
target_coef_df = readtable('../../Data/ModelData/TargetCoeffs.csv');
if readcog == 1
mean_var_df = readtable('../../Data/ModelData/occupation_exposed_task_shares_cognitive_3cat.csv');
end
if readcog == 0 
mean_var_df = readtable('../../Data/ModelData/occupation_exposed_task_shares_3cat.csv');
end 

if type == 1 %Make sure baseline calibration matches data coefficient closely 
% Rename and merge necessary columns
logl_df = OccRows(:, {'naics', 'soc', 'dlog_L_of'});
logz_df = FirmRows(:, {'naics', 'dlog_Z_f','dlog_alphaF'});

df1 = innerjoin(logl_df, logz_df, 'Keys', 'naics');
df1.soc = double(df1.soc);  % Ensure numeric for merge

df = innerjoin(df1, mean_var_df, 'Keys', 'soc');

label1 = 'Ratio 10yr';
label2 = 'Z 10yr';
reg_model = 'OLS';

% Extract target values
target_m1 = target_coef_df{strcmp(target_coef_df.label, label1), reg_model};
target_m2 = target_coef_df{strcmp(target_coef_df.label, label2), reg_model};

% Adjust for delta
df.mean = df.mean * delta;
df.variance = df.variance * delta^2;

% Construct regression variables
X = [ones(height(df),1), df.mean, df.variance, df.dlog_Z_f, df.dlog_alphaF];
y = df.dlog_L_of;

% Run OLS regression
b = regress(y, X);
beta1 = b(2); % Coefficient on mean
beta2 = b(3); % Coefficient on variance

ratio = beta1 / sqrt(beta2); 
targetratio = target_m1; 

disp("Model Ratio: " + num2str(ratio) )
disp("Target Ratio: " + num2str(targetratio))
end

end
