clear; clc;
set(0,'DefaultTextInterpreter','latex');
set(0,'DefaultLineLineWidth',1.0);
set(0,'DefaultAxesFontSize',15);
set(0,'DefaultTextFontSize',15);
set(0,'DefaultAxesLineWidth',1.0);

%% Parameters
alpha = 0.3;
beta  = 0.1;
kmax = 4;

alphasq = alpha^2;
kvals = linspace(0,kmax,400);
Nk = length(kvals);
nroots = 5;

omega = zeros(Nk,nroots);

%% Polynomial coefficient function
coeffs = @(k)[ ...
    1,...
    k,...
    -(4+(1+beta+alphasq)*k^2),...
    4*k - (1+beta+alphasq)*k^3,...
    beta*k^4 + 4*beta*k^2 + 3*alphasq*k^2,...
    beta*k^5 - 4*beta*k^3 - alphasq*k^3];

%% First k: just compute roots
omega(1,:) = roots(coeffs(kvals(1)));

%% Track branches
for i = 2:Nk
    
    r = roots(coeffs(kvals(i)));
    prev = omega(i-1,:);
    
    assigned = false(1,nroots);
    new_order = zeros(1,nroots);
    
    for j = 1:nroots
        
        dist = abs(r - prev(j));
        dist(assigned) = inf;
        
        [~,idx] = min(dist);
        new_order(j) = r(idx);
        assigned(idx) = true;
        
    end
    
    omega(i,:) = new_order;
    
end

omega_exact = kvals;

%% Plot real and imaginary parts with consistent colors

colors = lines(nroots);   % one color per root

% --- Real part ---
figure(1); clf;

tiledlayout(1,2);

nexttile
hold on
for j = 1:nroots
    plot(kvals, real(omega(:,j)), ...
        'Color', colors(j,:), 'LineWidth',1.8);
end
plot(kvals,omega_exact,'k','Linewidth',1.8);
plot(kvals,omega_exact*sqrt(beta),'--k','Linewidth',1);
text(2.2,2.47,'$\omega=kv_{\rm A}$','FontSize',14,'Interpreter','latex','Rotation',45)
text(2.3,2.3*sqrt(beta)+0.2,'$\omega=kc_{\rm s}$','FontSize',14,'Interpreter','latex','Rotation',atand(sqrt(beta)));
plot(kvals,2-kvals,'-.k','LineWidth',1);
text(2.3,-0.05,'$\omega=(2k_0-k)v_{\rm A}$','FontSize',14,'Interpreter','latex','Rotation',-45);
plot(1,1,'ok','MarkerSize',6,'MarkerFaceColor', 'k', 'MarkerEdgeColor', 'k');
plot(2/(1+sqrt(beta)),2*sqrt(beta)/(1+sqrt(beta)),'ok','MarkerSize',6,'MarkerFaceColor','w','MarkerEdgeColor','k');
xlabel('$k/k_0$')
ylabel('${\rm Re}(\omega/\omega_0)$')
axis equal
xticks([0 1 2 3 4])
axis([0 kmax -kmax kmax])

ax = gca;
ax.TickLabelInterpreter = 'latex';
grid on

% --- Imaginary part ---

nexttile
hold on
for j = 1:nroots
    plot(kvals, imag(omega(:,j)), ...
        'Color', colors(j,:), 'LineWidth',1.8);
end
xlabel('$k/k_0$')
ylabel('${\rm Im}(\omega/\omega_0)$')
axis equal
axis([0 kmax -kmax kmax])
ax = gca;
ax.TickLabelInterpreter = 'latex';
xticks([0 1 2 3 4])
grid on