% Eric LaForest, March 2010
% laforest AT eecg.utoronto.ca

function [results] = project2()
% Looking for chaos in driven FHN system

    function output = driven_FHN(t, input, A, omega) 
        V = input(1);
        W = input(2);
        X = input(3);
        Y = input(4);
        
        % Input Stimulus
        F = @(t) A * cos(omega * t);

        % Forced Negative Resistance Oscillator
        Ydot = 0.2 * (1 - X^2) * Y - X^3 + F(t);
        Xdot = Y;
        
        % Pre-set FitzHugh-Nagumo System driven by X
        Vdot = V - (V^3 / 3) - W + X;
        Wdot = 0.08 * (V + 0.7 - 0.8 * W);
        
        output = [Vdot; Wdot; Xdot; Ydot];
    end

    function  [T,R] = solve_driven_FHN(time_span, initial_conditions, A, omega)
        this_FHN = @(t,input) driven_FHN(t, input, A, omega);
        [T,R] = ode15s(this_FHN, time_span, initial_conditions);
    end



%
% Wolf's algorithm on analytic models
%

    function [time_span, deltaR1, R0_fiducial, R0_perturbed] = trajectories(time_span, deltaR0, system_function, R0_fiducial, R0_perturbed)
        % Guess based on log2(3) > 1 bit new information
        delta_target = deltaR0 * 3;
        
        time_step = time_span(end) - time_span(1);

        while(true)
            [T1_fiducial, R1_fiducial] = system_function(time_span, R0_fiducial);
            [T1_perturbed, R1_perturbed] = system_function(time_span, R0_perturbed);

            % Keep only the final values
            R1_fiducial = R1_fiducial(end,:);
            R1_perturbed = R1_perturbed(end,:);

            R0_fiducial = R1_fiducial;
            R0_perturbed = R1_perturbed;
            time_span = time_span + time_step;

            deltaR1 = R1_perturbed - R1_fiducial;

            for i=1:length(deltaR1)
                this_delta = deltaR1(i);
                if(abs(this_delta) > delta_target)
                    R0_perturbed(i) = R1_fiducial(i) + ( deltaR0 * sign(this_delta) );
                end
            end

            if(any(abs(deltaR1) > delta_target))
                break;
            end
        end
    end
    
    function [total_time, LLE] = converge(time_span, deltaR0, system_function, R0_fiducial, R0_perturbed)
        zero_template = zeros(size(R0_fiducial));

        LLE = zero_template; 
        LLE_previous = zero_template;
        LLE_accumulator = zero_template;
        
        convergence_epsilon =  0.01;
        convergence_count = 0;
        convergence_target = 10;
 
        total_time = 0;
        
        while(convergence_count < convergence_target)
            [time_span, deltaR1, R0_fiducial, R0_perturbed] = trajectories(time_span, deltaR0, system_function, R0_fiducial, R0_perturbed);

            lambda_local = log2(abs(deltaR1) ./ deltaR0);
            LLE_accumulator = LLE_accumulator + lambda_local;
            
            time_step = time_span(end) - time_span(1);
            total_time = total_time + time_step;
            
            LLE = LLE_accumulator ./ total_time;

            %{
            if(mod(total_time,100) == 0)
                disp(LLE);
            end
            %}
            
            divergence = abs(LLE - LLE_previous);
            LLE_previous = LLE;

            if (all(divergence < convergence_epsilon))
                convergence_count = convergence_count + 1;
            else
                convergence_count = 0;
            end
        end
    end
    
    function [total_time, LLE] = find_LLE_numeric(system_function, initial_conditions)
        % Should be small enough to avoid going through folding regions most of the time,
        % but large enough to have sufficient numerical accuracy.
        % Determined empirically as smallest interval that does not yield negative exponents, 
        % but needs verification.
        time_span = [0 8];

        % From Sprott (2003) Sec. 5.6
        % suitable for double-precision floats
        deltaR0 = 10^-10;
        
        R0_fiducial = initial_conditions;
        R0_perturbed = R0_fiducial + deltaR0;
        
        [total_time, LLE] = converge(time_span, deltaR0, system_function, R0_fiducial, R0_perturbed);
    end

    function [Vdot_LLE, Xdot_LLE] = find_LLE_Wolf_analytic(initial_conditions, A, omega)
        time_series_function = @(time_span, initial_conditions) solve_driven_FHN(time_span, initial_conditions, A, omega);
        [total_time, LLE] = find_LLE_numeric(time_series_function, initial_conditions);
        Vdot_LLE = LLE(1);
        Xdot_LLE = LLE(3);
    end



%
% Rosenstein's algorithm on experimentally-generated data
%

    function [points, count] = time_series_sizes(time_series)
       time_series_size = size(time_series);
       points = time_series_size(1);
       count = time_series_size(2);
    end

    function [time, delay_embedding] = delay_embed(time, time_series)
       lag = 1;
       embedding_dimension = 4;
       
       [time_series_points, time_series_count] = time_series_sizes(time_series);
       
       delay_embedding_points = time_series_points - (lag * embedding_dimension);
       time = time(1:delay_embedding_points);
       
       delay_embedding = zeros(delay_embedding_points, embedding_dimension, time_series_count);
       
       for i = 1:time_series_count
           for j = 1:delay_embedding_points
               for k = 1:embedding_dimension
                   delay_embedding(j,k,i) = time_series(j+k,i);
               end
           end
       end
    end

    function distance = neighbour_distance(A, B)
        distance = norm(A - B);
    end

    function mean_period = time_series_mean_period(time_series, end_time)
        [time_series_points, time_series_count] = time_series_sizes(time_series); 
        NFFT = 2^nextpow2(time_series_points);
        freq_domain = fft(time_series, NFFT) / time_series_points;
        mean_frequency = mean(2*abs(freq_domain(1:NFFT/2,:)));
        mean_period = (1 ./ mean_frequency) / (time_series_points / end_time);
    end

    function [index] = find_nearest_neighbour(time, delay_embeddings, mean_period, reference_point_index)
        % some huge distance never to be seen between two points in a series...
        distance = 1000;
        index = 0;
        %delta_t = 0;
        reference_point = delay_embeddings(reference_point_index);
        reference_time = time(reference_point_index);
        
        end_index = length(time);       
        for i = reference_point_index:end_index
            candidate = delay_embeddings(i);
            candidate_time = time(i);
            time_delta = candidate_time - reference_time;
            if(time_delta < mean_period)
                continue;
            end
            candidate_distance = neighbour_distance(reference_point, candidate);
            if(candidate_distance < distance)
                distance = candidate_distance;
                index = i;
                %delta_t = time_delta;
            end
        end
    end

    function [nearest_neighbours] = find_all_nearest_neighbours(time, delay_embedding, mean_period)
       neighbours_count = length(delay_embedding);
       nearest_neighbours = zeros(neighbours_count,1);
       for i = 1:neighbours_count
           index = find_nearest_neighbour(time, delay_embedding, mean_period, i);
           nearest_neighbours(i) = index;
       end
    end
    
    function [avg_dist] = find_avg_dist(delay_embedding, nearest_neighbour_indices, trajectory_step)
        distance_accumulator = 0;
        count = 0;
        last_index = length(delay_embedding);
        for i = 1:last_index
            reference_index = i + trajectory_step;
            if(reference_index > last_index)
                break;
            end
            reference = delay_embedding(reference_index);
            
            neighbour_index = nearest_neighbour_indices(i) + trajectory_step;
            if(neighbour_index > last_index)
                continue;
            end
            % points near the end have no nearest neighbour in the embedding
            % at a distance greater than the mean period
            if(neighbour_index == 0)
                continue;
            end
            neighbour = delay_embedding(neighbour_index);
            
            dist = log2(neighbour_distance(reference, neighbour));
            distance_accumulator = distance_accumulator + dist;
            count = count + 1;
        end
        avg_dist = distance_accumulator / count;
    end

    function [distances] = find_all_avg_dist(delay_embedding, nearest_neighbour_indices)
        % only the first few steps are valid, after that the lambda tends towards zero
        % as the trajectory folds over the attractor
        step_count = 5;
        for trajectory_step = 1:step_count
            distances(trajectory_step) = find_avg_dist(delay_embedding, nearest_neighbour_indices, trajectory_step);
        end
    end

    function [LLE] = fit_line(distances)
        x = [1:length(distances)];
        LLE = polyfit(x, distances, 1);
        % assume flat line if no fit possible (ie: LLE of zero)
        if(any(isnan(LLE)))
            LLE = [0 0];
        end
    end

    function [Vdot_LLE, Xdot_LLE] = find_LLE_Rosenstein_experimental(initial_conditions, A, omega)
        time_series_function = @(time_span, initial_conditions) solve_driven_FHN(time_span, initial_conditions, A, omega);
        delta_t = 0.1;
        [T,R] = time_series_function(0:delta_t:100, initial_conditions);

        [time, delay_embedding] = delay_embed(T, R);
        mean_period = time_series_mean_period(R, T(end));
    
        Vdot_embedding = delay_embedding(:,:,1);
        Vdot_mean_period = mean_period(1);
    
        Xdot_embedding = delay_embedding(:,:,3);
        Xdot_mean_period = mean_period(3);
    
        Vdot_nearest_neighbour_indices = find_all_nearest_neighbours(time, Vdot_embedding, Vdot_mean_period);
        Xdot_nearest_neighbour_indices = find_all_nearest_neighbours(time, Xdot_embedding, Xdot_mean_period);
    
        Vdot_distances = find_all_avg_dist(Vdot_embedding, Vdot_nearest_neighbour_indices);
        Xdot_distances = find_all_avg_dist(Xdot_embedding, Xdot_nearest_neighbour_indices);
    
        Vdot_line_coefficients = fit_line(Vdot_distances);
        Xdot_line_coefficients = fit_line(Xdot_distances);
        
        Vdot_LLE = Vdot_line_coefficients(1);
        Xdot_LLE = Xdot_line_coefficients(1);
    end


% Plotting

    % Because we deal with both large and small numbers
    format short g

    initial_conditions = [0,0,0,0];
    A = 17;     A_step = 0.1;    A_min = 16;     A_max = 18;
    omega = 4;  o_step = 0.01;    o_min = 2;      o_max = 5;

    A_all = A_min:A_step:A_max;
    omega_all = o_min:o_step:o_max;

    tic;
    points = length(A_all) * length(omega_all)
    % based on ~2sec per point on my machine
    estimated_seconds = (points*2)
    estimated_minutes = estimated_seconds / 60
    estimated_hours = estimated_minutes / 60
    results = zeros(points, 4);
    foo = 1;
    for Ai = A_all
        for omegai = omega_all
            [Vdot_LLE, Xdot_LLE] = find_LLE_Rosenstein_experimental(initial_conditions, Ai, omegai);
            results(foo,:) = [Ai, omegai, Vdot_LLE, Xdot_LLE];
            foo = foo + 1;
            percent_remaining = (points - foo) / points
        end
    end
    toc;
    save('rosenstein_results', 'results');
    
    subplot(4,1,1), plot(results(:,1));
    subplot(4,1,2), plot(results(:,2));
    subplot(4,1,3), plot(results(:,3));
    subplot(4,1,4), plot(results(:,4));
    
end