syms z;
syms t;
mass = 0.0001; % marble's mass
radius = 0.001; % marble's radius
V = (4/3)*pi*radius^3; % marble's volume
z0 = 0.03;
phi1 = 0; % phase at origin
c = 343; % celerity of sound in the air (20°C, m/s)
r = 1; % reflection coefficient; set to 1 (perfect mirror)
A = 0.00001; % amplitude of the incident wave (V ?)
l = 0.01; % position of the first node (m)
L = 2*l; % distance inter-mirror (m)
t_span = linspace(0.001,10*(2*L/c+0.001),1001);
t_step = t_span(2)-t_span(1);
k = 0.5*pi/l;
w = c*k;
rho = 1.2041; % density of air (20°C, kg/m^3)
n = 5; % number of nodes to display
zeros = linspace(l,(2*n+1)*pi/(2*k),n+1); % nodes positions

% phase shift for reflected wave is 0+phase of incident wave at z=L because of the acoustic impedance of the mirror > air
incident_wave = A*cos(k*z-w*t+phi1);
incident_wave_handle = @(z,t) eval(simplify(incident_wave));
reflected_wave = r*A*cos(-k*(z+L)-w*(t+L/c)+phi1);

% Total wave
f = simplify(incident_wave + reflected_wave);
f_handle = @(z,t) eval(f);

p = simplify((f^2)/(diff(f, t)) * w^2 * rho * c);
F_p = -diff(p,z);
F_p_handle = @(z,t) eval(simplify(F_p));

clear Z;
Z = 0*ones(1,length(t_span));
Z(1) = z0;
Z(2) = z0 + t_span(2)^2 *9.81*(-mass + rho*V);

for i = 1:length(t_span)-2
    Z(i+2) = t_span(i)^2 *(9.81*(-mass + rho*V) + subs(subs(F_p,z,Z(i)),t,t_span(i))*V)/mass + Z(i);
end

Z_table = smoothdata(Z',"gaussian","SmoothingFactor",0.25);
%plot(t_span,Z_table);

for i = 1:(length(zeros)-1)
    xline(zeros(i), '--', 'HandleVisibility', 'off'); % Vertical line at x_values(i) with red dashed line
end

% Add images
img1 = imread('radio_gen.png');
img2 = imread('mirror.png');

hold on;
image(img1, 'XData', [-0.5*L, 0], 'YData', 10*[-1.7*A, 1.7*A]);
image(img2,'XData', [zeros(end-1), zeros(end-1)+0.5*L], 'YData', 10*[-1.5*A, 1.5*A]);
hold off;
    
for tc = 1:length(t_span)
    tt = t_span(tc);

    yyaxis left;
    colororder([0, 0.5, 1]);
    fplot(f_handle(z, tt), [0,(2*n+1)*pi/(2*k)+1], 'LineWidth', 2, 'DisplayName', 'A(z)');
    ylabel('I(z); (V)');
    ylim([-2*A, 2*A]);
    hold on;
    fplot(incident_wave_handle(z, tt), [0,(2*n+1)*pi/(2*k)+1], '--',  'LineWidth', 1, 'DisplayName', 'A_i(z)')
    scatter(Z(round(tc)), 0, 'filled', 'MarkerFaceColor', 'black', 'DisplayName', 'marble');
    hold off;
    
    yyaxis right;
    colororder([1, 0.6, 0]);
    fplot(F_p_handle(z, tt)/1000000, [0,(2*n+1)*pi/(2*k)+1], 'LineWidth', 2, 'DisplayName', 'F_p(z)');
    ylabel('F_p(z); (N/cm^3)');
    ylim([-2*10, 2*10]);


    xlabel('z (m)');
    title(['t =' num2str(tt)]);

    xlim([-0.5*L, l*2*n-0.5*l]);
    legend;

    % update screen
    hold on;
    image(img1, 'XData', [-0.5*L, 0], 'YData', 10*[-2*A, 2*A]);
    image(img2,'XData', [zeros(end-1), zeros(end-1)+0.5*L], 'YData', 10*[-1.5*A, 1.5*A], AlphaData=0.8);
    hold off;
    drawnow
end