clc; clear; close all;
figure('units', 'normalized', 'outerposition', [0,0,1,1]);

N=51; % Number of balls

% Define the system as a first-order ODE system
dydt = @(t,y) [y(2); -y(1)*(y(1)^2 - 1)];

% Initial positions for 31 trajectories
x0_array = linspace(-1.75, 1.75, N);
v0 = 0;  % all start with zero velocity

% Time span
tspan = linspace(0, 20, 2000);

% Precompute trajectories
trajectories = zeros(length(tspan), length(x0_array));
velocities = zeros(length(tspan), length(x0_array));
for k = 1:length(x0_array)
    y0 = [x0_array(k); v0];
    [~, y] = ode23s(dydt, tspan, y0);
    trajectories(:,k) = y(:,1);   % store x(t) 
    velocities(:,k) = y(:,2);     % store dx/dt(t)
end

% Define potential V(x) = x^4/4 - x^2/2
V = @(x) 0.25*x.^4 - 0.5*x.^2;

colors = lines(length(x0_array));

%% Left: Potential plot
subplot(1,2,1);
hold on; grid on;
xrange = linspace(-2,2,500);
plot(xrange, V(xrange), 'k','LineWidth',2); % potential curve
ylim([-0.3 1]); 
xlabel('x','FontSize',16); ylabel('V(x)','FontSize',16);
title('Balls in Double-Well Potential','FontSize',16);

% Initialize balls
balls = gobjects(length(x0_array),1);
for k = 1:length(x0_array)
    balls(k) = plot(trajectories(1,k), V(trajectories(1,k)), 'o', ...
                    'MarkerSize', 8, 'MarkerFaceColor', colors(k,:), 'MarkerEdgeColor', colors(k,:));
end

%% Right: Phase space
subplot(1,2,2);
hold on; grid on;
xlabel('x','FontSize',16); ylabel('dx/dt','FontSize',16);
title('Phase Space','FontSize',16);
xlim([-2 2]); ylim([-2 2]);

% Initialize points and trajectory lines in phase space
phase_points = gobjects(length(x0_array),1);
phase_lines  = gobjects(length(x0_array),1);
for k = 1:length(x0_array)
    % Moving point
    phase_points(k) = plot(trajectories(1,k), velocities(1,k), 'o', ...
                           'MarkerSize', 6, 'MarkerFaceColor', colors(k,:), 'MarkerEdgeColor', colors(k,:));
    % Line showing trajectory so far
    phase_lines(k) = plot(trajectories(1, k), velocities(1, k), '-', ...
                          'Color', colors(k,:), 'LineWidth', 1);
end

%% Animate
for i = 1:length(tspan)
    % Update balls on potential
    subplot(1,2,1);
    for k = 1:length(x0_array)
        set(balls(k), 'XData', trajectories(i,k), 'YData', V(trajectories(i,k)));
    end
    
    % Update phase space points and lines
    subplot(1,2,2);
    for k = 1:length(x0_array)
        set(phase_points(k), 'XData', trajectories(i,k), 'YData', velocities(i,k));
        set(phase_lines(k), 'XData', trajectories(1:i,k), 'YData', velocities(1:i,k));
    end
    
    if mod(i,3)==1
        drawnow;
        if i==1
            pause
        end
    end
end