close all;
% Load data
data = load('traj-2.txt');

% Perform PCA and reduce the data
[eigenvectors, reduced_data] = reduce2(data);

% Number of clusters for discretization
num_clusters = 100;

% Discretize the reduced data
discretized_trajectory = discretize_trajectory(reduced_data, num_clusters);

% Calculate the count matrix
count_matrix = calculate_count_matrix(discretized_trajectory, num_clusters);

% Calculate the transition matrix
transition_matrix = calculate_transition_matrix(count_matrix);

% Calculate the eigenvalues and vectors ([V,D] = eig(A) returns diagonal matrix D of eigenvalues and)
[eigenvectors, eigenvaluesdiag] = eig(transition_matrix);

% eigenvaluesdiag is a diagonal matrix, so we need to extract the eigenvalues
eigenvalues = diag(eigenvaluesdiag);

% Omit the first eigenvectors, eigenvalues
eigenvectors = eigenvectors(:, 2:end);
eigenvalues = eigenvalues(2:end);

% Convert eigenvalues into timescales
lag_time = 1;
timescales = -lag_time ./ log(abs(eigenvalues));

% plot figure
figure;
imagesc(transition_matrix);
title('Transition Matrix');
xlabel('To State');
ylabel('From State');
colorbar;

figure;
bar(1:num_clusters, sum(count_matrix, 2), 'EdgeColor', 'none');
title('Count Matrix');
xlabel('State');
ylabel('Count');

figure;
plot(discretized_trajectory);
title('Discretized Trajectory');
xlabel('State');
ylabel('Timesteps');

figure;
scatter(data(:, 1),data(:, 2));
title('Scatter');
xlabel('X');
ylabel('Y');

% Plot the eigenfunctions using bar graph
figure;
plot(eigenvectors(:, 1));
title('Eigenfunctions');
xlabel('State');
ylabel('Eigenfunction Value');
grid on;

% Plot the timescales of secound, third, fourth eigenvalues
figure;
bar(2:5, timescales(1:4));
title('Timescales (2nd - 5th)');
xlabel('Eigenvalue');
ylabel('Timescale');
grid on;


function [eigenvectors, reduced_data] = reduce2(data)
    % Perform PCA on the trajectory data
    covariance_matrix = cov(data);
    [eigenvectors, ~] = eig(covariance_matrix);
    
    % Get the first principal component
    first_principal_component = eigenvectors(:, end);
    
    % Reduce data to 1 dimension
    reduced_data = data * first_principal_component;

    % Display reduced data
    disp('Reduced Data:');
    disp(reduced_data);
end

function discretized_trajectory = discretize_trajectory(trajectory, num_clusters)
    % Function to discretize the trajectory into clusters
    
    % Find the min and max value
    min_val = min(trajectory);
    max_val = max(trajectory);
    
    % Calculate bin width
    bin_width = (max_val - min_val) / num_clusters;

    % Discretize the trajectory
    discretized_trajectory = discretize(trajectory, min_val:bin_width:max_val);
end

function count_matrix = calculate_count_matrix(discretized_trajectory, num_clusters)
    % Function to calculate the count matrix from a discretized trajectory

    % Initialize count matrix
    count_matrix = zeros(num_clusters, num_clusters);

    % Iterate through the discretized trajectory to count transitions
    for i = 1:length(discretized_trajectory)-1
        from_state = discretized_trajectory(i);
        to_state = discretized_trajectory(i+1);
        count_matrix(from_state, to_state) = count_matrix(from_state, to_state) + 1;
    end
end

function transition_matrix = calculate_transition_matrix(count_matrix)
    % Function to calculate the transition matrix from the count matrix

    % Get the number of states
    num_clusters = size(count_matrix, 1);

    % Initialize transition matrix
    transition_matrix = zeros(num_clusters, num_clusters);

    % Calculate transition probabilities
    for i = 1:num_clusters
        total_transitions = sum(count_matrix(i, :));

        % Avoid division by zero
        if total_transitions > 0
            transition_matrix(i, :) = count_matrix(i, :) / total_transitions;
        end
    end
end
