function out = calibrate_zbar(alphaJ, alphaO, desired_w, desired_q, params, varargin)
% CALIBRATE_ZBAR  Solve for a single global zbar so a uniform (w*, q*) satisfies FOC.
%
% Inputs
%   alphaJ : cell {O x 1}, alphaJ{o} = [J_o x 1] task weights α_j for SOC o
%            (will be normalized inside so sum_j α_j = 1 for each SOC)
%   alphaO : [O x 1] occupation weights α_o (not required to sum to 1; GM branch normalizes)
%   desired_w : scalar baseline wage for ALL tasks (e.g., 1.0)
%   desired_q : scalar baseline capital price for ALL tasks (e.g., 1.0)
%   params : struct with fields (zbar solved here; ignored if present):
%       .beta, .zeta, .nu, .a, .b, .chi, .psi, .theta, .Ybar (unused here),
%       .M (optional; default 1.0).  <-- NEW: used consistently in p_j
%
% Optional name/value:
%   'Tol' : tolerance for 1/eta (GM) switches, default 1e-12
%
% Output
%   out : struct with
%       .zbar        : scalar normalizer so uniform (w*,q*) satisfies FOC
%       .Xof         : [O x 1] occupation aggregator X_{o,f} at (w*,q*)
%       .Pof         : [O x 1] price index P_{o,f} = 1 ./ X_{o,f}
%       .Zf          : scalar firm-level aggregator Z_f at (w*,q*)
%       .pj_per_occ  : cell {O x 1}, each [J_o x 1] p_j at (w*,q*)
%       .params_used : params echoed back (with defaults filled)
%
% Notes
%   - Under uniform w* and q*, the residual FOC is identical across tasks;
%     zbar is obtained in closed form from that FOC.
%   - This routine DOES NOT compute GE objects; it only pins down the PE
%     normalizer consistent with your per-firm FOC at the chosen (w*, q*).

% ---------------- Parse / defaults ----------------
p = inputParser;
addParameter(p, 'Tol', 1e-12);
parse(p, varargin{:});
tol = p.Results.Tol;

% Pull parameters with safe defaults
beta  = getd(params,'beta',  0.5);
zeta  = getd(params,'zeta',  0.3);
nu    = getd(params,'nu',    3.0);
a     = getd(params,'a',     1.0);
b     = getd(params,'b',     1.0);
chi   = getd(params,'chi',   0.5);
psi   = getd(params,'psi',   0.5);
theta = getd(params,'theta', 2.0);
M     = getd(params,'M',     1.0);   % <-- NEW: used in p_j

% Guards
if abs(nu-1)  < tol, error('calibrate_zbar: nu must not equal 1.');  end
if isempty(alphaJ) || ~iscell(alphaJ), error('alphaJ must be a {O x 1} cell.'); end
O = numel(alphaJ);
alphaO = alphaO(:);
if numel(alphaO) ~= O
    error('alphaO must have length equal to numel(alphaJ).');
end

% Normalize α_j within each occupation
alphaJn = alphaJ;
for o = 1:O
    aj = alphaJ{o}(:);
    s  = sum(aj);
    if s <= 0, aj = ones(size(aj)); s = numel(aj); end
    alphaJn{o} = aj / s;
end

% ---------------- Build uniform w*, q* arrays and aggregators ----------------
wstar = desired_w;
qstar = desired_q;

pj_per_occ = cell(O,1);
Xof        = zeros(O,1);
Pof        = zeros(O,1);

for o = 1:O
    J_o = numel(alphaJn{o});
    w   = wstar * ones(J_o,1);
    q   = qstar * ones(J_o,1);

    % p_j = [ a*(M*w)^(1-nu) + b*q^(1-nu) ]^(1/(1-nu))   <-- NEW: includes M
    pj = agg_task(w, q, a, b, nu);
    pj_per_occ{o} = pj;

    % X_{o,f} = [ sum_j α_j * p_j^(1-psi) ]^(-1/(1-psi))  (GM limit if psi ~ 1)
    Xof(o) = agg_occ(pj, alphaJn{o}, psi, tol);

    % P_{o,f} = X_{o,f}^{-1}
    Pof(o) = 1 ./ Xof(o);
end

% Firm-level aggregator (GM limit if chi ~ 1)
Zf = agg_firm(Pof, alphaO, chi, tol);

% ---------------- Closed-form zbar under uniform (w*, q*) ----------------
% From the task-level FOC in your solver, under uniform wages:
%   LHS = zbar * w^{1+zeta}
%   RHS = a*M^{-nu}*w^{1-nu} * ( a*M^{1-nu}*w^{1-nu} + b*q^{1-nu} )^{(nu-psi)/(1-nu)}
%         * X_{o,f}^{chi-psi} * Z_f^{theta-chi}
% and X_{o,f}, Z_f above are identical across o at uniform (w*,q*).
core_inner = a.*(M.^(1-nu)).*(wstar.^(1-nu)) + b.*(qstar.^(1-nu));
RHS_core   = a.*(M.^(-nu)).*(wstar.^(1-nu)) .* (core_inner.^((nu-psi)/(1-nu)));

% Choose any occupation’s X_{o,f}; they’re identical here
if O >= 1
    Xfac = (Xof(1)).^(chi - psi);
else
    Xfac = 1.0;
end
Zfac = Zf.^(theta - chi);

RHS_all = RHS_core .* Xfac .* Zfac;
LHS_fac = wstar.^(1 + zeta);

zbar = RHS_all ./ LHS_fac;

% ---------------- Pack outputs ----------------
out = struct();
out.zbar        = zbar;
out.Xof         = Xof;
out.Pof         = Pof;
out.Zf          = Zf;
out.pj_per_occ  = pj_per_occ;
out.params_used = struct('beta',beta,'zeta',zeta,'nu',nu,'a',a,'b',b, ...
                         'chi',chi,'psi',psi,'theta',theta,'M',M);

% ================ nested helpers =================
    function pj = agg_task(w, q, a_, b_, nu_)
        % Task price index with M included on wages
        t1 = a_ .* ((M .* w).^(1 - nu_));
        t2 = b_ .* (q.^(1 - nu_));
        pj = (t1 + t2).^(1/(1 - nu_));
    end

    function X = agg_occ(pj, aj, psi_, tol_)
        eta = 1 - psi_;
        if abs(eta) > tol_
            s = sum( aj(:) .* (pj(:).^(eta)) );
            X = s.^(-1/eta);
        else
            wts = aj(:) / sum(aj);
            X   = exp( -sum( wts .* log(pj(:)) ) );  % geometric-mean limit
        end
    end

    function Z = agg_firm(Pof_, ao, chi_, tol_)
        eta = 1 - chi_;
        if abs(eta) > tol_
            s = sum( ao(:) .* (Pof_(:).^(eta)) );
            Z = s.^(-1/eta);
        else
            wts = ao(:) / sum(ao);
            Z   = exp( -sum( wts .* log(Pof_(:)) ) ); % geometric-mean limit
        end
    end

end

% -------- small utility --------
function v = getd(s, k, dflt)
    if isfield(s,k) && ~isempty(s.(k)), v = s.(k); else, v = dflt; end
end
