classdef ODEgeodesics
    properties
        N                       % number of parameters
        M                       % number of dynamic variables

        theta                   % vector of symbols for parameters

        obs_config              % define observation configuration. Each row is one observation. First column indicates which variable observed. Second column is time point. 
        obs_data
        time_points_to_evaluate % time points that will be evaluated for making model predictions and derivatives

        x                       % vector of symbols for variables
        dx                      % vector of operatins defining the RHS of the ode model
        x0                      % vector of symbols for initial conditions 
        ODE_rhs                 
        ODE_init
        
        S
        dS
        S0
        ODE_1st_order_sensitivity_rhs
        ODE_1st_order_sensitivity_init
        
        SS
        dSS
        SS0
        ODE_2nd_order_sensitivity_rhs
        ODE_2nd_order_sensitivity_init
        
        para                  % para to start the geodesic path

        dtau = 0.001;
        max_integration_steps = 2000;
        integrate_method_opt = 'ODE_solver';                          % 'Euler' or 'ODE_solver'
        accerlation_compute_opt = '2nd_order_sensitivity_for_Gamma';  % '2nd_order_finite_difference' or '2nd_order_sensitivity_for_Gamma';
        smallest_singular_threshold = 1e-5;
        check_flipping_acceleration = 0;
        display_integration_process = 0;
        display_parameter_names = [];
    end
    
    methods
        function obj = ODEgeodesics(x, dx, x0, theta, obs_config)
            obj.x = x;
            obj.dx = dx;
            obj.x0 = x0;
            obj.theta = theta;
            obj.obs_config = obs_config;
            obj.time_points_to_evaluate = unique(obs_config(:,2));   
            if ~ismember(0,obj.time_points_to_evaluate)
                obj.time_points_to_evaluate = [0; obj.time_points_to_evaluate];   % when numerically integrate, evaluate at the observed time points. Important to add 0 as the first time point, because ODE45 views the first time point as the time of the initial condition
            end
            obj.M = length(x);              % number of dynamic variables
            obj.N = length(theta);          % number of parameters
            obj.ODE_rhs  = matlabFunction([dx],'vars', {'t', [x], theta});      % define the right hand side of the ODE as a function of time, system states + sensitivity, and parameter
            obj.ODE_init = matlabFunction([x0],'vars', {theta});                % define the initial condition as a function of parameters
            for k=1:obj.N
                obj.display_parameter_names{k,1} = char(theta(k));
            end
        end
        
        
        function obj = compute_1st_order_sensitivity_equations(obj)
            % compute first order sensitivity equations
            [obj.S,obj.dS,obj.S0] = sensitivity_equations(obj.x, obj.theta, obj.dx, obj.x0);
            obj.ODE_1st_order_sensitivity_rhs  = matlabFunction([obj.dx;obj.dS(:)],'vars', {'t', [obj.x;obj.S(:)], obj.theta});    % define the right hand side of the ODE as a function of time, system states + sensitivity, and parameter
            obj.ODE_1st_order_sensitivity_init = matlabFunction([obj.x0;obj.S0(:)],'vars', {obj.theta});                   % define the initial condition as a function of parameters
        end
        
        
        function obj = compute_2nd_order_sensitivity_equations(obj)
            M = obj.M;
            N = obj.N;
            if isempty(obj.S) || isempty(obj.dS) || isempty(obj.S0)
                obj = obj.compute_1st_order_sensitivity_equations;
            end
            % compute 2nd order sensitivity equations
            [SS,dSS,SS0] = sensitivity_equations([obj.x;obj.S(:)], obj.theta, [obj.dx;obj.dS(:)], [obj.x0;obj.S0(:)]); 
            SS(1:M,:)=[];  dSS(1:M,:)=[]; SS0(1:M,:)=[]; % remove the first order sensitivity equations
            obj.SS = reshape(SS,M,N,N);                      % reshape the 2nd order equations 
            obj.dSS = reshape(dSS,M,N,N);
            obj.SS0 = reshape(SS0,M,N,N);
            obj.ODE_2nd_order_sensitivity_rhs  = matlabFunction([obj.dx;obj.dS(:);obj.dSS(:)],'vars', {'t', [obj.x;obj.S(:);obj.SS(:)], obj.theta});    % define the right hand side of the ODE as a function of time, system states + sensitivity, and parameter
            obj.ODE_2nd_order_sensitivity_init = matlabFunction([obj.x0;obj.S0(:);obj.SS0(:)],'vars', {obj.theta});                                     % define the initial condition as a function of parameters
        end

        
        function model_predictions = compute_model_prediction(obj, para)
            if nargin==1
                para = obj.para;
            end
            [t,y] = ode45(@(t,x_values)obj.ODE_rhs(t,x_values,para),obj.time_points_to_evaluate,obj.ODE_init(para));  % numerically integrate, evaluate at the observed time points. Important to add 0 as the first time point, because ODE45 views the first time point as the time of the initial condition
            x_values = y(:,1:obj.M); y(:,1:obj.M) = [];   
            % pull out the simulated data corresponding to the observations
            for i=1:size(obj.obs_config,1)
                time_ind = ismember(t,obj.obs_config(i,2));
                var_ind  = obj.obs_config(i,1);
                % model prediction
                model_predictions(i,1) = x_values(time_ind, var_ind);
            end            
        end
        
        
        function [J,smallest_singular_vector,smallest_singular_value, singular_values] = compute_Jacobian(obj, para)
            if nargin==1
                para = obj.para;
            end
            [t,y] = ode45(@(t,x_values)obj.ODE_1st_order_sensitivity_rhs(t,x_values,para),obj.time_points_to_evaluate,obj.ODE_1st_order_sensitivity_init(para));  
            x_values = y(:,1:obj.M); y(:,1:obj.M) = [];   
            x_1st_order = y(:,1:obj.M*obj.N); y(:,1:obj.M*obj.N)=[];
            % pull out the simulated data corresponding to the observations
            for i=1:size(obj.obs_config,1)
                time_ind = ismember(t,obj.obs_config(i,2));
                var_ind  = obj.obs_config(i,1);
                % model prediction
                model_predictions(i,1) = x_values(time_ind, var_ind);
                % jacobian
                tmp = reshape(x_1st_order(time_ind,:),obj.M,obj.N);
                J(i,:) = tmp(var_ind,:);
            end
            if sum(isnan(J(:)))~=0 || sum(isinf(J(:)))~=0
                singular_values = zeros(min(obj.N,obj.M),1);
                smallest_singular_value = 0;
                smallest_singular_vector = zeros(obj.N,1);
            else
                [U, singular_values, V] = svd(J); 
                if size(singular_values,1)==1 || size(singular_values,2)==1
                    singular_values = singular_values(1,1);
                else
                    singular_values = diag(singular_values);
                end
                [smallest_singular_value,ind] = min(singular_values); 
                smallest_singular_vector = V(:,ind);
            end
        end
        
        
        function [Gamma, J] = compute_Gamma_connection(obj, para)
            if nargin==1
                para = obj.para;
            end
            [t,y] = ode45(@(t,x_values)obj.ODE_2nd_order_sensitivity_rhs(t,x_values,para),obj.time_points_to_evaluate,obj.ODE_2nd_order_sensitivity_init(para));  
            % separate the numerically simulated data into the system itself, 1st order, 2nd order sentitivities 
            x_values = y(:,1:obj.M); y(:,1:obj.M) = [];   
            x_1st_order = y(:,1:obj.M*obj.N); y(:,1:obj.M*obj.N)=[];
            x_2nd_order = y; 
            % pull out the simulated data corresponding to the observations
            for i=1:size(obj.obs_config,1)
                time_ind = ismember(t,obj.obs_config(i,2));
                var_ind  = obj.obs_config(i,1);
                % model prediction
                model_predictions(i,1) = x_values(time_ind, var_ind);
                % jacobian
                tmp = reshape(x_1st_order(time_ind,:),obj.M,obj.N);
                J(i,:) = tmp(var_ind,:);
                % 2nd order derivatives
                tmp = reshape(x_2nd_order(time_ind,:),obj.M,obj.N,obj.N);
                H(i,:,:) = tmp(var_ind,:,:);
            end
            g = J'*J; % fisher information matrix
            g_inv = inv(g);
            Gamma = reshape(g_inv*(J'*reshape(H,size(obj.obs_config,1),obj.N*obj.N)),obj.N,obj.N,obj.N);  
        end
        

        function [velocity_acceleration] = compute_geodesic_ODE_RHS(obj, t, para_and_velocity)
            velocity_acceleration = zeros(obj.N*2,1);
            para = para_and_velocity(1:obj.N);
            velocity = para_and_velocity(obj.N+1:end);
            velocity_acceleration(1:obj.N) = velocity;
            
            if isequal(obj.accerlation_compute_opt, '2nd_order_finite_difference')  % Transtrum, PLoS CB 2016, eqn (55-57)
                h = 1e-4;
                Avv = ( compute_model_prediction(obj, para+h*velocity) + compute_model_prediction(obj, para-h*velocity) - 2*compute_model_prediction(obj, para) ) /h/h;
                [J,smallest_singular_vector,smallest_singular_value, singular_values] = compute_Jacobian(obj, para);
                if sum(isnan(J(:)))~=0 || sum(isinf(J(:)))~=0
                    acceleration = zeros(obj.N,1);
                else
                    acceleration = -pinv(J)*Avv;
                end
            elseif isequal(obj.accerlation_compute_opt, '2nd_order_sensitivity_for_Gamma')
                [Gamma, ~] = compute_Gamma_connection(obj, para);
                for mu=1:obj.N
                    tmp = 0;
                    for alfa = 1:obj.N
                        for beta = 1:obj.N
                             tmp = tmp + Gamma(mu,alfa,beta)*velocity(alfa,1)*velocity(beta,1);
                        end
                    end
                    acceleration(mu,1) = -tmp;
                end
            else
                error('incorrect option set for the method for computing acceleration!')
            end
            velocity_acceleration(obj.N+1:end) = acceleration;
            velocity_acceleration(isnan(velocity_acceleration))=0;
        end

        
        function [para_direction, para_acceleration, singular_values] = compute_initial_velocity(obj, para)
            if nargin==1
                para = obj.para;
            end
            [J,smallest_singular_vector,smallest_singular_value, singular_values] = compute_Jacobian(obj, para);
            para_direction = smallest_singular_vector;
            [tmp_rhs] = compute_geodesic_ODE_RHS(obj, 0, [para;para_direction]);
            para_acceleration = tmp_rhs(obj.N+1:end);
            if para_direction'*para_acceleration<0
                para_direction = -para_direction;
            end
        end
        
        
        function [geodesic_path, tau] = numerical_integrate_for_geodesic_path(obj, para)
            % para : the parameter value to start the geodesic path. If not provide, the default is the para property in obj 
            if nargin<=1
                para = obj.para;
            end

            
            if isequal(obj.integrate_method_opt,'Euler')  % integrate the geodesic path by the simple Euler's method
                dtau = obj.dtau;
                max_integration_steps = obj.max_integration_steps;
                [para_direction, para_acceleration, singular_values] = compute_initial_velocity(obj, para);
                y = [para', para_direction']; 
                tau = 0;
                for i=1:max_integration_steps
                    [tmp_rhs] = compute_geodesic_ODE_RHS(obj, 0, [y(i,:)']);
                    y(i+1,:) = y(i,:)+tmp_rhs'*dtau;
                    tau(i+1,1) = tau(i,1) + dtau;
                    current_para = y(i,1:obj.N)';
                    current_direction = y(i,obj.N+1:end)';
                    current_acceleration = tmp_rhs(obj.N+1:end);
                    [J,smallest_singular_vector,smallest_singular_value, singular_values] = compute_Jacobian(obj, current_para);
                    if smallest_singular_value<obj.smallest_singular_threshold
                        display('may need to stop, FIM singular'); 
                        y(end,:) = [];
                        tau(end) = [];
                        break;
                    end
                    if obj.check_flipping_acceleration==1
                        if current_direction'*current_acceleration<0
                            display('may need to stop, acceleration direction flip'); 
                            y(end-1:end,:) = [];
                            tau(end-1:end) = [];
                            break;
                        end
                    end
                    if obj.display_integration_process==1
                        disp(['current tau: ', num2str(tau(i)), ' - current smallest singular value: ', num2str(smallest_singular_value)])
                        %figure(999)
                        %subplot(1,2,1); plot(y(:,1:3),'linewidth',3); legend(obj.display_parameter_names)
                        %subplot(1,2,2); plot(y(:,4:6),'linewidth',3); legend(strcat('Velocity_{', obj.display_parameter_names, '}'))
                        %title(['current smallest singular value: ', num2str(smallest_singular_value)])
                        %drawnow
                    end
                end
                geodesic_path = y;

            elseif  isequal(obj.integrate_method_opt,'ODE_solver') % the other option is 'ODE_solver'
                
                dtau = obj.dtau;
                max_integration_steps = obj.max_integration_steps;
                tau = 0:dtau:dtau*max_integration_steps;
                [para_direction, para_acceleration, singular_values] = compute_initial_velocity(obj, para);
                opt = odeset('Events',@(tau,y) obj.stopping_event_function(tau, y));
                [tau,y] = ode15s(@(t,y_values)obj.compute_geodesic_ODE_RHS(t,y_values),...
                              [tau],[para;para_direction], opt);  % numerically integrate, evaluate at the observed time points. Important to add 0 as the first time point, because ODE45 views the first time point as the time of the initial condition
                geodesic_path = y;
                
            else
                
                error('incorrect option chosen for the integration method!')
                
            end
        end
        
        
        function [value,isterminal,direction] = stopping_event_function(obj, tau, y)
            % when value is equal to zero, an event is triggered.
            % set isterminal to 1 to stop the solver at the first event, or 0 to
            % get all the events.
            %  direction=0 if all zeros are to be computed (the default), +1 if
            %  only zeros where the event function is increasing, and -1 if only
            %  zeros where the event function is decreasing.

            y = y(:)';
            [tmp_rhs] = compute_geodesic_ODE_RHS(obj,0,[y']);
            current_para = y(1:obj.N)';
            current_direction = y(obj.N+1:2*obj.N)';
            current_acceleration = tmp_rhs(obj.N+1:2*obj.N);
            [J,smallest_singular_vector,smallest_singular_value, singular_values] = compute_Jacobian(obj, current_para);
            if obj.display_integration_process==1
                disp(['current tau: ', num2str(tau), ' - current smallest singular value: ', num2str(smallest_singular_value)])
            end
            if (obj.check_flipping_acceleration==1 && current_direction'*current_acceleration<0) || smallest_singular_value<obj.smallest_singular_threshold || norm(current_acceleration)==0
                if (obj.check_flipping_acceleration==1 && current_direction'*current_acceleration<0)
                    display('may need to stop, acceleration direction flip'); 
                end
                if smallest_singular_value<obj.smallest_singular_threshold
                    display('may need to stop, FIM singular'); 
                end
                if norm(current_acceleration)==0
                    display('may need to stop, acceleration all NaN or all 0'); 
                end
                value = 0;  % when value = 0, an event is triggered
                isterminal = 1; % terminate after the first event
                direction = 0;  % get all the zeros
            else
                value = 1;  % when value = 0, an event is triggered
                isterminal = 1; % terminate after the first event
                direction = 0;  % get all the zeros
            end
        end

    end
end



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [S,dS,S0] = sensitivity_equations(x, theta, dx, x0)
    M=length(x);        % number of dynamic variables
    N=length(theta);    % number of parameters

    % first order sensitivity
    for i=1:M
        for j=1:N
            S(i,j) = sym(['d',char(x(i)),'_d',char(theta(j))]); %sym(['S[',num2str(i),[','],num2str(j),']']);
        end
    end

    for i=1:M
        for j=1:N
            dS(i,j) = diff(dx(i),theta(j),1);
            for l = 1:M
                dS(i,j) = dS(i,j) + diff(dx(i), x(l), 1) * S(l,j);
            end
            S0(i,j) = diff(x0(i), theta(j), 1);
        end
    end
end
