function results = vibfatigue_optimizer_eband(opts)
%% VIBFATIGUE_OPTIMIZER_EBAND  (v5-compatible)
%
% Optimizer using E_band (dominant-mode energy fraction) as the smooth
% correction scale factor for the Modified Dirlik model.
%
% CORRECTION MODEL:
%
%   D_meta = D_DKI * exp( (1-E_band) * [b0*(1-a2^2) + b1*a2*(1-a2)] )
%
%   alpha2  = m2/sqrt(m0*m4)          spectral bandwidth
%   E_band  = energy in octave around dominant peak / total energy
%             stored in shapeOut(:,16) from vibfatigue_psd_mc_v5
%
% V5 ADDITIONS:
%   - 3D stratified split: alpha2 x E_band x psdType
%   - psdType-aware diagnostics (Type 1=pure peaks, 2=peaks+floor,
%     3=shaped broadband) stored in vf_meta_inputs.psdTypes
%   - subplot grid fixed to 2x5 to accommodate 9 k-values (k=3:11)
%   - Per-psdType LogRMSE reported for all three models
%
% FITS THREE MODELS for comparison:
%   Model 1 — Raw DKI (baseline, no correction)
%   Model 2 — 2-feature (alpha2 only, no E_band)
%   Model 3 — 3-feature (alpha2 + E_band smooth scale)  <- primary model
%
% REQUIRES: vf_meta_inputs from vibfatigue_psd_mc_v5
%
% USAGE:
%   results = vibfatigue_optimizer_eband()
%   results = vibfatigue_optimizer_eband(opts)
%
% opts fields (all optional):
%   opts.trainFrac    - train fraction          (default 0.80)
%   opts.rngSeed      - split seed              (default 42)
%   opts.lambda       - Tikhonov weight         (default 0.001)
%   opts.momentRelTol - relative moment floor   (default 1e-6)
%   opts.verbose      - print tables            (default true)
%   opts.plotDiag     - plot diagnostics        (default true)
%   opts.maxFunEvals  - fminsearch evals        (default 80000)
%   opts.nRestarts    - multi-start restarts    (default 5)

if nargin < 1, opts = struct; end
if ~isfield(opts,'trainFrac'),    opts.trainFrac    = 0.80;   end
if ~isfield(opts,'rngSeed'),      opts.rngSeed      = 42;     end
if ~isfield(opts,'lambda'),       opts.lambda       = 0.001;  end
if ~isfield(opts,'momentRelTol'), opts.momentRelTol = 1e-6;   end
if ~isfield(opts,'verbose'),      opts.verbose      = true;   end
if ~isfield(opts,'plotDiag'),     opts.plotDiag     = true;   end
if ~isfield(opts,'maxFunEvals'),  opts.maxFunEvals  = 80000;  end
if ~isfield(opts,'nRestarts'),    opts.nRestarts    = 5;      end

fprintf('\n========================================\n');
fprintf('  VibFatigue Optimizer — E_band Model (v5)\n');
fprintf('========================================\n\n');

%% ---- 1. Load data --------------------------------------------------
try
    vf = evalin('base', 'vf_meta_inputs');
catch
    error('vf_meta_inputs not found. Run vibfatigue_psd_mc_v5 first.');
end

kvals  = vf.kvals;
n_k    = numel(kvals);
Ncases = size(vf.momentsOut, 1);
fprintf('Dataset: %d cases, %d k-values\n', Ncases, n_k);

% Check E_band is present (v4/v5 MC required)
if size(vf.shapeOut, 2) < 16 || all(vf.shapeOut(:,16) == 0)
    error(['shapeOut column 16 (E_band) is missing or all zeros.\n' ...
           'Re-run the MC with vibfatigue_psd_mc_v5.']);
end

% PSD type (v5: col 10 of diagOut; 1=pure peaks, 2=peaks+floor, 3=broadband)
if isfield(vf, 'psdTypes') && ~all(vf.psdTypes == 0)
    psdTypes = vf.psdTypes;
    has_psdTypes = true;
    fprintf('v5 data detected: psdType breakdown — ');
    for t = 1:3
        fprintf('Type%d=%d  ', t, sum(psdTypes==t));
    end
    fprintf('\n');
else
    psdTypes = ones(Ncases, 1);   % treat all as Type 1 if not present
    has_psdTypes = false;
    fprintf('[Info] psdTypes not found — treating all cases as Type 1\n');
end
fprintf('\n');

%% ---- 2. Spectral parameters ----------------------------------------
m0     = vf.momentsOut(:,2);
m2     = vf.momentsOut(:,4);
m4     = vf.momentsOut(:,5);
alpha2 = max(0, min(1, m2 ./ (sqrt(m0.*m4) + 1e-9)));
E_band = max(0, min(1, vf.shapeOut(:,16)));
nModes = vf.shapeOut(:,15);

tol   = opts.momentRelTol;
valid = (m0 > tol*max(m0)) & (m2 > tol*max(m2)) & (m4 > tol*max(m4));
fprintf('Valid cases: %d / %d\n\n', sum(valid), Ncases);

% Distribution summary
fprintf('--- Dataset statistics (valid cases) ---\n');
fprintf('alpha2 : min=%.3f  p25=%.3f  median=%.3f  p75=%.3f  max=%.3f\n', ...
        min(alpha2(valid)), prctile(alpha2(valid),25), median(alpha2(valid)), ...
        prctile(alpha2(valid),75), max(alpha2(valid)));
fprintf('E_band : min=%.3f  p25=%.3f  median=%.3f  p75=%.3f  max=%.3f\n', ...
        min(E_band(valid)), prctile(E_band(valid),25), median(E_band(valid)), ...
        prctile(E_band(valid),75), max(E_band(valid)));
fprintf('nModes : ');
for n = 0:5
    cnt = sum(nModes(valid)==n);
    if cnt > 0, fprintf('%d=%d  ', n, cnt); end
end
fprintf('\n\n');

%% ---- 3. Method index -----------------------------------------------
names  = vf.methodNames;
idxDKI = find(strcmp(names,'DKI'));
assert(~isempty(idxDKI), 'DKI not found in methodNames.');

%% ---- 4. 3D stratified train/test split (alpha2 x E_band x psdType) ----
% v5 addition: include psdType as a stratification dimension so all three
% PSD topology classes are proportionally represented in train and test.
rng(opts.rngSeed);
idx_all = find(valid);

a2_edges = quantile(alpha2(idx_all), [0 1/3 2/3 1]);
eb_edges = quantile(E_band(idx_all), [0 1/3 2/3 1]);
a2_edges(1) = 0;   a2_edges(end) = 1.001;
eb_edges(1) = 0;   eb_edges(end) = 1.001;

ptype_vals = unique(psdTypes(idx_all));   % typically [1 2 3]

idx_train = [];
idx_test  = [];
for ia = 1:3
    for ie = 1:3
        for ip = 1:numel(ptype_vals)
            pt = ptype_vals(ip);
            in_s = idx_all( alpha2(idx_all)   >= a2_edges(ia)   & ...
                            alpha2(idx_all)   <  a2_edges(ia+1) & ...
                            E_band(idx_all)   >= eb_edges(ie)   & ...
                            E_band(idx_all)   <  eb_edges(ie+1) & ...
                            psdTypes(idx_all) == pt );
            n_s = numel(in_s);
            if n_s == 0, continue; end
            n_tr_s = round(opts.trainFrac * n_s);
            perm_s = in_s(randperm(n_s));
            idx_train = [idx_train; perm_s(1:n_tr_s)];       %#ok
            idx_test  = [idx_test;  perm_s(n_tr_s+1:end)];   %#ok
        end
    end
end
n_train = numel(idx_train);
n_test  = numel(idx_test);
fprintf('3D stratified split (alpha2 x E_band x psdType): Train=%d  Test=%d\n', ...
        n_train, n_test);

% Report psdType distribution in train/test
if has_psdTypes
    fprintf('  Train psdType: ');
    for t = 1:3
        fprintf('T%d=%d(%.0f%%)  ', t, sum(psdTypes(idx_train)==t), ...
                100*mean(psdTypes(idx_train)==t));
    end
    fprintf('\n  Test  psdType: ');
    for t = 1:3
        fprintf('T%d=%d(%.0f%%)  ', t, sum(psdTypes(idx_test)==t), ...
                100*mean(psdTypes(idx_test)==t));
    end
    fprintf('\n');
end
fprintf('\n');

%% ---- 5. Detect v2/v4 vs v5 data ------------------------------------
is_rate = isfield(vf,'T_actual');
if ~is_rate
    fprintf('[Info] v2 data: dividing by Tdur=%.0fs\n\n', vf.opts.Tdur);
end

%% ---- 6. Cache per-k data -------------------------------------------
a_tr   = alpha2(idx_train);
a_te   = alpha2(idx_test);
eb_tr  = E_band(idx_train);
eb_te  = E_band(idx_test);
a_all  = alpha2;
eb_all = E_band;

logC_train  = zeros(n_train, n_k);
Ddki_all    = zeros(Ncases,  n_k);
Drf_all     = zeros(Ncases,  n_k);

for j = 1:n_k
    k_val        = kvals(j);
    Drf          = get_Drf(vf, j, is_rate);
    Dall         = get_Dall(vf, k_val, is_rate);
    Ddki         = Dall(:, idxDKI);
    Ddki_all(:,j) = Ddki;
    Drf_all(:,j)  = Drf;
    logC_train(:,j) = log(max(1e-20, Drf(idx_train))) - ...
                      log(max(1e-20, Ddki(idx_train)));
end

%% ---- 7. Tikhonov second-difference matrix --------------------------
D2 = zeros(n_k-2, n_k);
for ii = 1:n_k-2
    D2(ii,ii)=1; D2(ii,ii+1)=-2; D2(ii,ii+2)=1;
end

fs_opts = optimset('MaxFunEvals', opts.maxFunEvals, ...
                   'MaxIter',     opts.maxFunEvals, ...
                   'TolFun',      1e-11, ...
                   'TolX',        1e-11, ...
                   'Display',     'off');

%% ================================================================
%%  MODEL 2: 2-feature reference (alpha2 only)
%% ================================================================
fprintf('--- Model 2: 2-feature reference (alpha2 only) ---\n');

X2_tr  = basis2(a_tr);
X2_all = basis2(a_all);

B2_ols = zeros(n_k, 2);
for j = 1:n_k
    B2_ols(j,:) = (X2_tr \ logC_train(:,j)).';
end

obj2 = @(b) tik_obj(b, X2_tr, logC_train, D2, opts.lambda, n_k, 2);
[b2_opt, ~] = fminsearch(obj2, B2_ols(:), fs_opts);
B2 = reshape(b2_opt, n_k, 2);

err_dki_tr = zeros(1,n_k); err_dki_te = zeros(1,n_k);
err2_tr    = zeros(1,n_k); err2_te    = zeros(1,n_k);
for j = 1:n_k
    [err_dki_tr(j), err_dki_te(j), err2_tr(j), err2_te(j)] = ...
        eval_linear(B2(j,:), basis2(a_tr), basis2(a_te), ...
                    Ddki_all(:,j), Drf_all(:,j), idx_train, idx_test);
end
fprintf('  DKI   : train=%.4f  test=%.4f\n', mean(err_dki_tr), mean(err_dki_te));
fprintf('  2-feat: train=%.4f  test=%.4f  (improv vs DKI: %.1f%%)\n\n', ...
        mean(err2_tr), mean(err2_te), 100*(1-mean(err2_te)/mean(err_dki_te)));

%% ================================================================
%%  MODEL 3: E_band smooth scale (primary model)
%%
%%  log_C = (1-E_band) * [b0*(1-a2^2) + b1*a2*(1-a2)]
%% ================================================================
fprintf('--- Model 3: E_band smooth scale ---\n');
fprintf('    log_C = (1-E_band)*[b0*(1-a2^2) + b1*a2*(1-a2)]\n\n');

X3_tr  = basis3(a_tr,  eb_tr);
X3_te  = basis3(a_te,  eb_te);
X3_all = basis3(a_all, eb_all);

B3_ols = zeros(n_k, 2);
for j = 1:n_k
    B3_ols(j,:) = (X3_tr \ logC_train(:,j)).';
end

fprintf('  OLS coefficients [b0, b1] per k:\n');
for j = 1:n_k
    fprintf('    k=%-3d  b0=%.4f  b1=%.4f\n', kvals(j), B3_ols(j,1), B3_ols(j,2));
end
fprintf('\n');

obj3 = @(b) tik_obj(b, X3_tr, logC_train, D2, opts.lambda, n_k, 2);

best_fval = inf;
best_b3   = B3_ols(:);

for r = 1:opts.nRestarts
    if r == 1
        b_init = B3_ols(:);
    else
        b_init = B3_ols(:) + 0.1*randn(size(B3_ols(:)));
    end
    [b_r, fval_r] = fminsearch(obj3, b_init, fs_opts);
    if fval_r < best_fval
        best_fval = fval_r;
        best_b3   = b_r;
    end
    fprintf('  Restart %d/%d: fval=%.6f\n', r, opts.nRestarts, fval_r);
end
fprintf('  Best objective = %.6f\n\n', best_fval);

B3 = reshape(best_b3, n_k, 2);

err3_tr = zeros(1,n_k); err3_te = zeros(1,n_k);

if opts.verbose
    fprintf('  %-5s  %-10s %-10s %-10s %-10s  %-8s %-8s\n', ...
            'k','DKI_te','2feat_te','3feat_te','Improv%','b0','b1');
    fprintf('  %s\n', repmat('-',1,68));
end

for j = 1:n_k
    [~, ~, err3_tr(j), err3_te(j)] = ...
        eval_linear(B3(j,:), X3_tr, X3_te, ...
                    Ddki_all(:,j), Drf_all(:,j), idx_train, idx_test);
    if opts.verbose
        fprintf('  k=%-3d  %-10.4f %-10.4f %-10.4f %-10.1f  %-8.4f %-8.4f\n', ...
                kvals(j), err_dki_te(j), err2_te(j), err3_te(j), ...
                100*(1-err3_te(j)/err_dki_te(j)), B3(j,1), B3(j,2));
    end
end

fprintf('\n  Mean LogRMSE (all types):\n');
fprintf('    Raw DKI : test=%.4f\n', mean(err_dki_te));
fprintf('    2-feat  : test=%.4f  (%+.1f%% vs DKI)\n', ...
        mean(err2_te), 100*(mean(err2_te)/mean(err_dki_te)-1));
fprintf('    3-feat  : test=%.4f  (%+.1f%% vs DKI)\n', ...
        mean(err3_te), 100*(mean(err3_te)/mean(err_dki_te)-1));
fprintf('    3-feat vs 2-feat: %+.1f%%\n\n', ...
        100*(mean(err3_te)/mean(err2_te)-1));

%% ---- Per-psdType breakdown (v5 addition) ---------------------------
if has_psdTypes
    type_labels = {'Type1 (pure peaks)', 'Type2 (peaks+floor)', 'Type3 (broadband)'};
    fprintf('  Per-psdType LogRMSE (test set, mean over k):\n');
    fprintf('  %-22s  %-10s %-10s %-10s\n', 'PSD Type','DKI','2-feat','3-feat');
    fprintf('  %s\n', repmat('-',1,56));
    for t = 1:3
        idx_t = idx_test(psdTypes(idx_test) == t);
        if isempty(idx_t), continue; end
        dki_t = zeros(1,n_k); m2_t = zeros(1,n_k); m3_t = zeros(1,n_k);
        for j = 1:n_k
            dki_t(j) = logrms(Ddki_all(idx_t,j), Drf_all(idx_t,j));
            D2mod     = Ddki_all(idx_t,j) .* exp(basis2(alpha2(idx_t)) * B2(j,:).');
            D3mod     = Ddki_all(idx_t,j) .* exp(basis3(alpha2(idx_t), E_band(idx_t)) * B3(j,:).');
            m2_t(j)  = logrms(D2mod, Drf_all(idx_t,j));
            m3_t(j)  = logrms(D3mod, Drf_all(idx_t,j));
        end
        fprintf('  %-22s  %-10.4f %-10.4f %-10.4f\n', ...
                type_labels{t}, mean(dki_t), mean(m2_t), mean(m3_t));
    end
    fprintf('\n');
end

%% ---- Validation on real-world case ---------------------------------
fprintf('--- Validation: real case (a2=0.2047, E_band=0.974, k=6.54) ---\n');
a2_real    = 0.2047;
Eband_real = 0.974;
k_real     = 6.54;
D_DKI_real = 1.7021e-9;
D_rf_real  = 1.7920e-9;

b2_r = zeros(1,2); b3_r = zeros(1,2);
for c = 1:2
    b2_r(c) = interp1(kvals, B2(:,c), k_real, 'pchip');
    b3_r(c) = interp1(kvals, B3(:,c), k_real, 'pchip');
end

logC_2feat = basis2(a2_real)             * b2_r.';
logC_3feat = basis3(a2_real, Eband_real) * b3_r.';

D_2feat = D_DKI_real * exp(logC_2feat);
D_3feat = D_DKI_real * exp(logC_3feat);

scale_real = 1 - Eband_real;

fprintf('  Ground truth (rainflow) : %.4e /s\n', D_rf_real);
fprintf('  Raw Dirlik              : %.4e /s  (error %+.1f%%)\n', ...
        D_DKI_real, 100*(D_DKI_real/D_rf_real-1));
fprintf('  2-feat (alpha2 only)    : %.4e /s  (error %+.1f%%)\n', ...
        D_2feat, 100*(D_2feat/D_rf_real-1));
fprintf('  3-feat (E_band scale)   : %.4e /s  (error %+.1f%%)\n', ...
        D_3feat, 100*(D_3feat/D_rf_real-1));
fprintf('\n  E_band scale (1-E_band) : %.4f\n', scale_real);
fprintf('  log_C (2-feat)          : %.4f  (corr=%.4f)\n', logC_2feat, exp(logC_2feat));
fprintf('  log_C (3-feat)          : %.4f  (corr=%.4f)\n', logC_3feat, exp(logC_3feat));

%% ---- Coefficient table for implementation --------------------------
fprintf('\n--- Coefficient table for vibfatigue_modified_dirlik.m ---\n');
fprintf('k_cal = ['); fprintf('%d ', kvals); fprintf('];\n');
fprintf('B_cal = [\n');
fprintf('%% b0          b1\n');
for j = 1:n_k
    fprintf('    %10.6f,  %10.6f;  %% k=%d\n', B3(j,1), B3(j,2), kvals(j));
end
fprintf('];\n\n');

%% ---- Diagnostic plots ----------------------------------------------
if opts.plotDiag
    plot_eband_diagnostics(B2, B3, alpha2, E_band, psdTypes, has_psdTypes, ...
                           kvals, idx_train, idx_test, Ddki_all, Drf_all, ...
                           err_dki_te, err2_te, err3_te, ...
                           a2_real, Eband_real, D_rf_real, D_DKI_real);
end

%% ---- Pack outputs --------------------------------------------------
results.B2          = B2;
results.B3          = B3;
results.kvals       = kvals;
results.err_dki_te  = err_dki_te;
results.err2_te     = err2_te;
results.err3_te     = err3_te;
results.idx_train   = idx_train;
results.idx_test    = idx_test;
results.alpha2      = alpha2;
results.E_band      = E_band;
results.psdTypes    = psdTypes;

assignin('base','vf_eband_results', results);
fprintf('vf_eband_results assigned to base workspace.\n');

end % ---- end main ----


%% ========================================================================
%%                      BASIS FUNCTIONS
%% ========================================================================

function X = basis2(a2)
X = [(1-a2(:).^2),  a2(:).*(1-a2(:))];
end

function X = basis3(a2, Eband)
scale = 1 - Eband(:);
X = [scale.*(1-a2(:).^2),  scale.*a2(:).*(1-a2(:))];
end


%% ========================================================================
%%                      OBJECTIVE & EVALUATION
%% ========================================================================

function err = tik_obj(b, X_tr, logC_train, D2, lambda, n_k, n_coeff)
B          = reshape(b, n_k, n_coeff);
pred       = X_tr * B.';
resid      = pred - logC_train;
ols_err    = sum(resid(:).^2);
smooth_err = 0;
for p = 1:n_coeff
    smooth_err = smooth_err + sum((D2*B(:,p)).^2);
end
err = ols_err + lambda * smooth_err;
end

function [dki_tr, dki_te, mod_tr, mod_te] = eval_linear( ...
        b, X_tr, X_te, Ddki_all, Drf_all, idx_train, idx_test)
Dmod_tr = Ddki_all(idx_train) .* exp(X_tr * b(:));
Dmod_te = Ddki_all(idx_test)  .* exp(X_te * b(:));
dki_tr  = logrms(Ddki_all(idx_train), Drf_all(idx_train));
dki_te  = logrms(Ddki_all(idx_test),  Drf_all(idx_test));
mod_tr  = logrms(Dmod_tr,             Drf_all(idx_train));
mod_te  = logrms(Dmod_te,             Drf_all(idx_test));
end

function lr = logrms(D_pred, D_ref)
eps_d = 1e-20;
lr    = sqrt(mean((log10(max(eps_d,D_pred))-log10(max(eps_d,D_ref))).^2));
end


%% ========================================================================
%%                      DATA HELPERS
%% ========================================================================

function Drf = get_Drf(vf, j, is_rate)
if is_rate,  Drf = vf.DrfOut(:,j+1);
else,        Drf = vf.DrfOut(:,j+1) / vf.opts.Tdur;
end
end

function Dall = get_Dall(vf, k_val, is_rate)
Dall = vf.(sprintf('Dspec_k%d', k_val));
if ~is_rate, Dall = Dall / vf.opts.Tdur; end
end


%% ========================================================================
%%                      DIAGNOSTIC PLOTS
%% ========================================================================

function plot_eband_diagnostics(B2, B3, alpha2, E_band, psdTypes, has_psdTypes, ...
                                 kvals, idx_train, idx_test, Ddki_all, Drf_all, ...
                                 err_dki_te, err2_te, err3_te, ...
                                 a2_real, Eband_real, D_rf_real, D_DKI_real)
n_k    = numel(kvals);
a_plot = linspace(0,1,150).';

% Marker colours per psdType for scatter plots
type_colors = [0.2 0.4 0.8;   % Type1 pure peaks — blue
               0.1 0.7 0.3;   % Type2 peaks+floor — green
               0.9 0.3 0.1];  % Type3 broadband — red

%% Figure 1: Correction surface slices — 2x5 grid for 9 k-values
% Slots 1-9 = k-values, slot 10 = coefficient plot
Eb_slices = [0.10, 0.50, 0.97];
lc_colors = {'b','g','r'};
lc_labels = {'E_{band}=0.10 (flat broadband)', ...
             'E_{band}=0.50 (mixed)', ...
             'E_{band}=0.97 (single mode, real case)'};

figure('Name','E_band Correction Surface','Position',[50 50 1800 720]);
for j = 1:n_k
    C_data = log(max(1e-20,Drf_all(idx_train,j))) - ...
             log(max(1e-20,Ddki_all(idx_train,j)));

    subplot(2,5,j);

    if has_psdTypes
        % Colour scatter by psdType
        for t = 1:3
            mask = psdTypes(idx_train) == t;
            if any(mask)
                scatter(alpha2(idx_train(mask)), C_data(mask), 6, ...
                        type_colors(t,:), 'filled', 'MarkerFaceAlpha', 0.4);
                hold on;
            end
        end
    else
        scatter(alpha2(idx_train), C_data, 8, E_band(idx_train), 'filled');
        colormap(gca, 'cool'); clim([0 1]); hold on;
    end

    for qi = 1:3
        lc_sl = basis3(a_plot, Eb_slices(qi)*ones(size(a_plot))) * B3(j,:).';
        plot(a_plot, lc_sl, lc_colors{qi}, 'LineWidth', 2, ...
             'DisplayName', lc_labels{qi});
    end

    lc_2f = basis2(a_plot) * B2(j,:).';
    plot(a_plot, lc_2f, 'k--', 'LineWidth', 1.2, 'DisplayName', '2-feat ref');

    lc_real = basis3(a2_real, Eband_real) * B3(j,:).';
    plot(a2_real, lc_real, 'kp', 'MarkerSize',12, 'LineWidth',2, ...
         'DisplayName','Real case');

    yline(0,'k:');
    xlabel('\alpha_2'); ylabel('log(D_{rf}/D_{DKI})');
    title(sprintf('k=%d | DKI %.3f \\rightarrow 3f %.3f', ...
                  kvals(j), err_dki_te(j), err3_te(j)));
    ylim([-0.5 0.5]);
    if j==1
        if has_psdTypes
            legend({'T1-peaks','T2-pk+fl','T3-broad', lc_labels{:}, '2-feat ref','Real case'}, ...
                   'Location','northeast','FontSize',6);
        else
            legend('Location','northeast','FontSize',6);
        end
    end
    hold off;
end

% Slot 10: coefficient plot
subplot(2,5,10);
hold on;
plot(kvals, B3(:,1), 'b-o', 'LineWidth',2, 'DisplayName','b_0');
plot(kvals, B3(:,2), 'r-s', 'LineWidth',2, 'DisplayName','b_1');
yline(0,'k--'); xlabel('k'); ylabel('Coefficient');
title('Coefficients vs k'); legend('Location','best','FontSize',8); grid on;

if has_psdTypes
    type_str = 'Blue=T1(peaks) Green=T2(pk+fl) Red=T3(broad)';
else
    type_str = 'Colour = E_{band}: blue=broadband, red=single mode';
end
sgtitle({'Model 3: log_C = (1-E_{band})[b_0(1-\alpha_2^2) + b_1\alpha_2(1-\alpha_2)]', ...
         type_str}, 'FontSize',10);

%% Figure 2: Model comparison + E_band effect
figure('Name','E_band Analysis','Position',[100 100 1300 480]);

subplot(1,3,1);
hold on;
plot(kvals, err_dki_te, 'k--o', 'LineWidth',1.5, 'DisplayName','Raw DKI');
plot(kvals, err2_te,    'b-s',  'LineWidth',2,   'DisplayName','2-feat (a2)');
plot(kvals, err3_te,    'r-^',  'LineWidth',2,   'DisplayName','3-feat (E_band)');
xlabel('k'); ylabel('Test LogRMSE');
title('Model comparison across k');
legend('Location','best'); grid on;

subplot(1,3,2);
if has_psdTypes
    for t = 1:3
        mask = valid_idx_for_scatter(psdTypes, t);
        scatter(alpha2(mask), E_band(mask), 6, type_colors(t,:), 'filled', ...
                'MarkerFaceAlpha', 0.3, 'DisplayName', ...
                sprintf('Type%d', t));
        hold on;
    end
    legend('Location','best','FontSize',8);
else
    scatter(alpha2, E_band, 6, 'b', 'filled');
    hold on;
end
plot(a2_real, Eband_real, 'kp', 'MarkerSize',14, 'LineWidth',2, ...
     'DisplayName','Real case');
xlabel('\alpha_2'); ylabel('E_{band}');
title('E_{band} vs \alpha_2 (by psdType)');
grid on;

subplot(1,3,3);
j6   = find(kvals >= 6, 1);
eb_range = linspace(0,1,200).';
lc_curve = basis3(a2_real*ones(200,1), eb_range) * B3(j6,:).';
lc_2feat = basis2(a2_real) * B2(j6,:).';

hold on;
plot(eb_range, exp(lc_curve), 'r-', 'LineWidth',2, 'DisplayName','3-feat');
yline(exp(lc_2feat), 'b--', 'LineWidth',1.5, 'Label','2-feat');
yline(1.0,           'k:',  'LineWidth',1.0, 'Label','Raw DKI');
xline(Eband_real,    'k--', 'LineWidth',1.5);

lc_r = basis3(a2_real, Eband_real) * B3(j6,:).';
D_r  = D_DKI_real * exp(lc_r);
plot(Eband_real, exp(lc_r), 'rp', 'MarkerSize',12, 'LineWidth',2, ...
     'DisplayName', sprintf('Real (err=%+.1f%%)', 100*(D_r/D_rf_real-1)));

xlabel('E_{band}'); ylabel('Correction factor');
title(sprintf('Correction vs E_{band} at \\alpha_2=%.2f, k=%d', ...
              a2_real, kvals(j6)));
legend('Location','best','FontSize',8); grid on; ylim([0.9 1.5]);

sgtitle('E_{band} Feature Analysis');

%% Figure 3: Parity plots at k=5, coloured by psdType (v5 addition)
j5 = find(kvals==5, 1); if isempty(j5), j5=3; end
Drf_te   = Drf_all(idx_test, j5);
Ddki_te  = Ddki_all(idx_test, j5);
D2_te    = Ddki_te .* exp(basis2(alpha2(idx_test)) * B2(j5,:).');
D3_te    = Ddki_te .* exp(basis3(alpha2(idx_test), E_band(idx_test)) * B3(j5,:).');

figure('Name','Parity plots k=5','Position',[150 150 1200 420]);
models = {Ddki_te, D2_te, D3_te};
titles = {'Raw Dirlik', '2-feat (alpha2)', '3-feat (E_band)'};
for p = 1:3
    subplot(1,3,p);
    if has_psdTypes
        for t = 1:3
            mask_t = psdTypes(idx_test) == t;
            loglog(Drf_te(mask_t), models{p}(mask_t), '.', ...
                   'Color', type_colors(t,:), 'MarkerSize', 5);
            hold on;
        end
        if p==1
            legend('Type1','Type2','Type3','Location','northwest','FontSize',7);
        end
    else
        loglog(Drf_te, models{p}, 'b.', 'MarkerSize', 4);
        hold on;
    end
    lims = [min([Drf_te; models{p}]), max([Drf_te; models{p}])];
    plot(lims, lims,     'k-',  'LineWidth',1.5);
    plot(lims, lims*1.5, 'r--', 'LineWidth',1.0);
    plot(lims, lims/1.5, 'r--', 'LineWidth',1.0);
    xlabel('D_{rainflow}'); ylabel('D_{predicted}');
    title(sprintf('%s\nLogRMSE=%.4f', titles{p}, logrms(models{p},Drf_te)));
    grid on; axis equal; hold off;
end
sgtitle('Parity plots — k=5, test set (dashed=±50%, colour=psdType)');

%% Figure 4 (v5 new): Per-psdType LogRMSE vs k
if has_psdTypes
    figure('Name','Per-psdType Performance','Position',[200 200 900 380]);
    type_labels = {'Type1: pure peaks', 'Type2: peaks+floor', 'Type3: broadband'};
    ls_dki  = {'k--o','k--s','k--^'};
    ls_3f   = {'b-o', 'g-s', 'r-^'};
    hold on;
    for t = 1:3
        idx_t = idx_test(psdTypes(idx_test) == t);
        if isempty(idx_t), continue; end
        dki_t = zeros(1,n_k); m3_t = zeros(1,n_k);
        for j = 1:n_k
            dki_t(j) = logrms(Ddki_all(idx_t,j), Drf_all(idx_t,j));
            D3mod     = Ddki_all(idx_t,j) .* exp(basis3(alpha2(idx_t), E_band(idx_t)) * B3(j,:).');
            m3_t(j)  = logrms(D3mod, Drf_all(idx_t,j));
        end
        plot(kvals, dki_t, ls_dki{t}, 'LineWidth',1.2, ...
             'DisplayName', ['DKI ' type_labels{t}]);
        plot(kvals, m3_t,  ls_3f{t},  'LineWidth',2.0, ...
             'DisplayName', ['3-feat ' type_labels{t}]);
    end
    xlabel('k'); ylabel('Test LogRMSE');
    title('Per-psdType: Raw DKI (dashed) vs 3-feat Model (solid)');
    legend('Location','best','FontSize',7); grid on;
end

end % plot_eband_diagnostics


%% ---- Helper: indices for scatter by type ---------------------------
function idx = valid_idx_for_scatter(psdTypes, t)
idx = find(psdTypes == t);
end
