rng(0,'twister')

% Values of r on which we are going to evaluate the correlation dimension
rs = [28:5:50]; %linspace(5,100,100) % CHANGE ME
sigma = 10;
b = 8.0/3.0;

% Increasing these increases accuracy, but kills performances. 
n_points = 50000; % Number of points in each trajectory
n_reject = 5000; % Reject the n_reject first point as transient. 
max_time = 100000; % Final time of each trajectory
n_average = 1000;  % Number of points on which the correlation dimension (averaged pointwise) is computed. 

% For each point in n_average, should we use the same trajectory (and just
% change the point from which we evaluate the pointwise dimension), or also
% redraw a trajectory ? 
redraw_traj = false; % In the class, we rather used =false; 


% Distance (epsilon) bounds for the pointwise and correlation dimension
min_dist = 0.01;
max_dist = 70.00;
n_epsilon = 100;

% Epsilons at which we are going to evaluate N(epsilon)
epsilon   = logspace(log10(min_dist), log10(max_dist), n_epsilon); 

% Epsilon bounds used to estimate the correlation dimension (simple linear
% fit over this restricted range).
eps_min = 0.1; 
eps_max = 10; 
correlation_dimension = NaN(size(rs));

for l = 1:numel(rs)
    r = rs(l); 
    disp(l);

    if(~redraw_traj)
        y0 = [5.*rand(1,1); 5.*rand(1,1); 5.*rand(1,1)];
        F = @(t,y) [sigma*(y(2)-y(1));  r*y(1)-y(2)-y(1)*y(3); y(1)*y(2)-b*y(3)];
        tspan = linspace(0, max_time, n_points);
        [t,y] = ode45(F,tspan,y0);
        % We discard the first n_reject points, to keep only the long term dynamics
        y = y((n_reject+1):end,:);
    end
    
    % We now evaluate the correlation dimension (in a naive way)
    % This is done by picking randomly n_average points over the attractor (for
    % performance reason. One could pick more, or even all of them). 
    
    N_epsilon = zeros(size(epsilon));
    for counter = 1:n_average
        if(redraw_traj)
            y0 = [5.*rand(1,1); 5.*rand(1,1); 5.*rand(1,1)];
            F = @(t,y) [sigma*(y(2)-y(1));  r*y(1)-y(2)-y(1)*y(3); y(1)*y(2)-b*y(3)];
            tspan = linspace(0, max_time, n_points);
            [t,y] = ode45(F,tspan,y0);
            % We discard the first n_reject points, to keep only the long term dynamics
            y = y((n_reject+1):end,:);
        end

        chosen_point = randi([1, size(y,1)], 1,1);
        X = y(chosen_point,:);
    
        % distances from X 
        dist = sqrt((y(:,1) - X(1)).^2 + (y(:,2) - X(2)).^2 + (y(:,3) - X(3)).^2);
        dist(chosen_point) = [];
            
        for k=1:numel(epsilon)
            N_epsilon(k) = N_epsilon(k) + sum(dist <= epsilon(k));
        end
    end
    N_epsilon = N_epsilon./n_average;
    % We renormalize just for the plotting
    figure(1); loglog(epsilon, N_epsilon./max(N_epsilon)); hold on; drawnow;
    xlabel('\epsilon')
    ylabel('N(\epsilon)')

    %% Fit over a very restricted range 
    i = epsilon > eps_min & epsilon < eps_max & N_epsilon > 0;
    epsilon_fit = epsilon(i);
    N_epsilon_fit = N_epsilon(i); 
    f=fit(log10(epsilon_fit)',log10(N_epsilon_fit)','a*x+b');
    
    %% Output the correlation dimension
    correlation_dimension(l) = f.a;
end

figure(2);
plot(rs, correlation_dimension, 'ko-');
xlabel('r');
ylabel('Correlation dimension');