function [W, Zf, Z_component, details] = solve_task_wages(Q, alphaJ, alphaO, alphaf, params, w0)
%SOLVE_TASK_WAGES  Jointly solve task wages {w_j} across occupations within a firm.
%
% Inputs
%   Q       : cell {O x 1}, Q{o} is [J_o x 1] vector of q_j for occupation o
%   alphaJ  : cell {O x 1}, alphaJ{o} is [J_o x 1] vector of α_j for occupation o
%   alphaO  : [O x 1] vector of α_{o,f}
%   alphaf  : scalar, industry f weight 
%   params  : struct with fields (defaults shown below if missing):
%               .beta=0.5, .zeta=0.3, .zbar=1.0, .nu=3.0, .a=1.0, .b=1.0, ...
%               .chi=0.5, .psi=0.5, .theta=2.0,
%               .M=1.0 (NEW safe default), .rhs_scale=1.0 (NEW safe default)
%   w0      : (optional) cell {O x 1}, initial guesses for w_j; if omitted, uses q_j
%
% Outputs
%   W            : cell {O x 1}, solved wages W{o} is [J_o x 1]
%   Zf           : scalar firm-level aggregator Z_f
%   Z_component  : scalar, Z_f^(theta - chi)
%   details      : struct with fields:
%                    .Xof (cell), .Pof ([O x 1]), .pj_cell (cell), .Zf (scalar)
%
% Notes
%   - Solves in log-wage space for positivity.
%   - Uses geometric-mean limits when ψ or χ approach 1.

    % ---------------- Parameters & guards ----------------
    beta  = getd(params,'beta',  0.5);
    zeta  = getd(params,'zeta',  0.3);
    zbar  = getd(params,'zbar',  1.0);
    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);

    % --- NEW optional flags (safe defaults) ---
    use_fixed_Zf = getd(params,'use_fixed_Zf',false);
    Zf_fixed     = getd(params,'Zf_fixed',[]);

    % --- NEW: safe defaults used by GE wrapper ---
    M         = getd(params,'M',         1.0);
    rhs_scale = getd(params,'rhs_scale', 1.0);

    tol = 1e-12;   % for ν, ψ, χ ~ 1 cases

    if abs(nu-1)  < tol, error('solve_task_wages: nu must not equal 1.');  end
    if abs(psi-1) < tol && abs(psi-1) >= 1e-16
        % we'll use geometric-mean limit internally; no error
    end
    if abs(chi-1) < tol && abs(chi-1) >= 1e-16
        % we'll use geometric-mean limit internally; no error
    end

    % ---------------- Shapes & initial guess ----------------
    O = numel(Q);
    if numel(alphaJ) ~= O || numel(alphaO) ~= O
        error('solve_task_wages: size mismatch between Q, alphaJ, and alphaO.');
    end

    J = cellfun(@numel, Q(:));
    idx_start = cumsum([1; J(1:end-1)]);
    idx_end   = cumsum(J);
    Jtot      = sum(J);

    % initial guess in log space
    if nargin < 5 || isempty(w0)
        w0 = cell(O,1);
        for o = 1:O
            % positive, benign initial guess
            w0{o} = max(Q{o}, 1);
        end
    end
    x0 = zeros(Jtot,1);
    for o = 1:O
        ii = idx_start(o):idx_end(o);
        x0(ii) = log(max(w0{o}(:), 1e-6));
    end

    % ---------------- Solve residuals = 0 in log-wage space ----------------
    fun = @(x) residuals_logw(x, Q, alphaJ, alphaO, ...
                              struct('beta',beta,'zeta',zeta,'zbar',zbar,'nu',nu, ...
                                     'a',a,'b',b,'chi',chi,'psi',psi,'theta',theta, ...
                                     'M',M,'rhs_scale',rhs_scale), ...
                              idx_start, idx_end);

    use_fsolve = (exist('fsolve','file') == 2);
    if use_fsolve
        optsFS = optimoptions('fsolve','Display','off','FunctionTolerance',1e-10, ...
                              'StepTolerance',1e-10, 'MaxIterations', 400, 'MaxFunctionEvaluations', 1e5);
        [x, ~, exitflag] = fsolve(fun, x0, optsFS);
        if exitflag <= 0
            warning('solve_task_wages: fsolve did not converge, attempting a few damped steps.');
            x = damped_steps(fun, x0, 60);
        end
    else
        % small damped fallback
        x = damped_steps(fun, x0, 100);
    end

    % ---------------- Unpack solution & compute objects ----------------
    W = cell(O,1);
    for o = 1:O
        ii = idx_start(o):idx_end(o);
        W{o} = exp(x(ii));
    end

    % Recompute pj, Xof, Pof, Zf for outputs
    pj_cell = cell(O,1);
    Xof     = cell(O,1);
    Pof     = zeros(O,1);

    for o = 1:O
        w = W{o}(:);
        q = Q{o}(:);

        % p_j = [ a*(M*w)^(1-nu) + b*q^(1-nu) ]^(1/(1-nu))
        t1 = a .* ((M.*w).^(1-nu));
        t2 = b .* (q.^(1-nu));
        pj = (t1 + t2).^(1/(1-nu));
        pj_cell{o} = pj;

        % X_{o,f} = [ sum_j α_j * p_j^(1-psi) ]^(-1/(1-psi)) (GM limit if psi ~ 1)
        if abs(psi-1) > tol
            s = sum( alphaJ{o}(:) .* (pj.^(1-psi)) );
            Xof{o} = s.^(-1/(1-psi));
        else
            wts   = alphaJ{o}(:) / sum(alphaJ{o}(:));
            Xof{o}= exp( -sum( wts .* log(pj) ) );
        end

        Pof(o) = 1 ./ Xof{o};
    end
    

    % Z_f = [ sum_o α_o * P_{o,f}^{1-chi} ]^(-1/(1-chi)) (GM limit if chi ~ 1)
    % Firm-level aggregator Z_f (or freeze to provided value)
    if use_fixed_Zf && ~isempty(Zf_fixed)
        Zf = Zf_fixed;   % <-- NEW: hold industry-level aggregator fixed
    else
    if abs(chi-1) > tol
        sZ = sum( alphaO(:) .* (Pof(:).^(1-chi)) );
        Zf = sZ.^(-1/(1-chi));
    else
        wts = alphaO(:) / sum(alphaO);
        Zf  = exp( -sum( wts .* log(Pof(:)) ) );
    end
    end

    Z_component = Zf^(theta - chi);

    % Details for GE aggregation & diagnostics
    details = struct();
    details.Xof     = Xof;
    details.Pof     = Pof;      % NEW: exposed
    details.pj_cell = pj_cell;

    % Z_f used for outputs (respect the freeze flag)
    if use_fixed_Zf && ~isempty(Zf_fixed)
        Zf_used = Zf_fixed;    % <-- NEW
    else
        if abs(chi-1) > tol
            sZ = sum( alphaO(:) .* (Pof(:) .^(1-chi)) );
            Zf_used = sZ.^(-1/(1-chi));
    else
        wts = alphaO(:) / sum(alphaO);
        Zf_used = exp( -sum( wts .* log(Pof(:)) ) );
        end
    end

    Zf = Zf_used;              % keep return value consistent
    details.Zf = Zf_used;      % already returned; now it’s the “used” value
    details.Zf      = Zf;       % NEW: exposed

    % === NEW: compute Y_f from (l_j,k_j) all the way up (primal) ===
    try
        % Normalize α_j per occupation if not already
        alphaJ_norm = alphaJ;
        for o=1:numel(alphaJ_norm), aj = alphaJ_norm{o}(:); alphaJ_norm{o} = aj / sum(aj); end
    
        [Yf_exact, dY] = Yf_from_wq(W, Q, alphaJ_norm, alphaO, alphaf, params);
        details.Yf   = Yf_exact;    % firm output (primal)
        details.Yof  = dY.Yof;      % per-occupation outputs
        details.lj   = dY.lj_cell;  % task labor
        details.kj   = dY.kj_cell;  % task capital
        details.yj   = dY.yj_cell;  % task outputs
    catch ME
        warning('Yf_from_wq failed in solve_task_wages: %s. Using dual fallback (Zf).', ME.message);
        if ~isfield(details,'Zf') || isempty(details.Zf)
            details.Zf = Zf;  % whatever you already computed
        end
        details.Yf = details.Zf;     % growth-safe fallback
    end

end % main

% =====================================================================
% Residuals in log-wage space (positivity guaranteed)
% =====================================================================
function [r, Xof, Pof, Zf, pj_cell] = residuals_logw(x, Q, alphaJ, alphaO, params, idx_start, idx_end)
    beta  = params.beta;
    zeta  = params.zeta;
    zbar  = params.zbar;
    nu    = params.nu;
    a     = params.a;
    b     = params.b;
    chi   = params.chi;
    psi   = params.psi;
    theta = params.theta;

    % --- NEW: safe defaults again inside residual (for robustness) ---
    if ~isfield(params,'M') || isempty(params.M), params.M = 1.0; end
    if ~isfield(params,'rhs_scale') || isempty(params.rhs_scale), params.rhs_scale = 1.0; end
    M         = params.M;
    rhs_scale = params.rhs_scale;

    tol = 1e-12;

    O = numel(Q);
    Jtot = idx_end(end);
    r = zeros(Jtot,1);

    % Unpack wages
    W = cell(O,1);
    for o = 1:O
        ii = idx_start(o):idx_end(o);
        W{o} = exp(x(ii));
    end

    % Precompute pj, Xof, Pof
    pj_cell = cell(O,1);
    Xof     = cell(O,1);
    Pof     = zeros(O,1);

    for o = 1:O
        w = W{o}(:);
        q = Q{o}(:);

        % p_j = [ a*(M*w)^(1-nu) + b*q^(1-nu) ]^(1/(1-nu))
        t1 = a .* ((M.*w).^(1-nu));
        t2 = b .* (q.^(1-nu));
        pj = (t1 + t2).^(1/(1-nu));
        pj_cell{o} = pj;

        % X_{o,f}
        if abs(psi-1) > tol
            s = sum( alphaJ{o}(:) .* (pj.^(1-psi)) );
            Xof{o} = s.^(-1/(1-psi));
        else
            wts = alphaJ{o}(:) / sum(alphaJ{o}(:));
            Xof{o} = exp( -sum( wts .* log(pj) ) ); % geometric-mean limit
        end

        Pof(o) = 1 ./ Xof{o};
    end

    % Firm-level aggregator Z_f
    if abs(chi-1) > tol
        sZ = sum( alphaO(:) .* (Pof(:).^(1-chi)) );
        Zf = sZ.^(-1/(1-chi));
    else
        wts = alphaO(:) / sum(alphaO);
        Zf  = exp( -sum( wts .* log(Pof(:)) ) );
    end

    % Residuals: per-task FOC
    for o = 1:O
        ii = idx_start(o):idx_end(o);
        w  = W{o};
        q  = Q{o}(:);

        % LHS: w_j^(1/beta) * (sum_j α_j w_j^(1/beta))^(beta - 1 + zeta*beta) * zbar
        sumw_alpha = sum( alphaJ{o}(:) .* (w.^(1/beta)) );
        LHS_factor = (sumw_alpha)^(beta - 1 + zeta*beta) * zbar;
        LHS = (w.^(1/beta)) .* LHS_factor;

        % RHS core: a*M^{-nu}*w^{1-nu} * ( a*M^{1-nu}*w^{1-nu} + b*q^{1-nu} )^{(nu-psi)/(1-nu)}
        core_inner = a.*(M.^(1-nu)).*(w.^(1-nu)) + b.*(q.^(1-nu));
        RHS_core   = a.*(M.^(-nu)).*(w.^(1-nu)) .* (core_inner.^((nu-psi)/(1-nu)));

        % Multiply by X_{o,f}^{chi-psi} * Z_f^{theta-chi}
        Xfac = (Xof{o}).^(chi - psi);
        Zfac = Zf.^(theta - chi);

        % --- NEW: scale RHS by params.rhs_scale (for GE blue factor gY) ---
        RHS = (RHS_core .* Xfac .* Zfac) .* rhs_scale;

        r(ii) = LHS - RHS;
    end
end

% =====================================================================
% Small damped fallback if fsolve is unavailable
% =====================================================================
function x = damped_steps(fun, x0, iters)
    x = x0;
    for it = 1:iters
        r = fun(x);
        nr = max(abs(r));
        if nr < 1e-10, break; end
        % crude diagonal preconditioner: scale by local magnitude
        step = -r ./ max(1,abs(r));
        lam  = 0.3;
        improved = false;
        for ls = 1:8
            x_try = x + lam*step;
            if max(abs(fun(x_try))) < nr
                x = x_try; improved = true; break;
            end
            lam = lam/2;
        end
        if ~improved
            x = x + 1e-3*step;
        end
    end
end

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