classdef ExplicitRungeKuttaDiscreteData < time.Timestepper
    properties
        D
        S           % Function handle for time-dependent data
        data        % Matrix of data vectors, one column per stage
        k
        t
        v
        m
        n
        order
        a, b, c, s  % Butcher tableau
        K           % Stage rates
        U           % Stage approximations
        T           % Stage times
    end


    methods
        function obj = ExplicitRungeKuttaDiscreteData(D, S, data, k, t0, v0, order)
            default_arg('order', 4);
            default_arg('S', []);
            default_arg('data', []);

            obj.D = D;
            obj.S = S;
            obj.k = k;
            obj.t = t0;
            obj.v = v0;
            obj.m = length(v0);
            obj.n = 0;
            obj.order = order;
            obj.data = data;

            switch order
            case 4
                [obj.a, obj.b, obj.c, obj.s] = time.rkparameters.rk4();
            otherwise
                error('That RK method is not available');
            end

            obj.K = sparse(obj.m, obj.s);
            obj.U = sparse(obj.m, obj.s);

        end

        function [v,t,U,T,K] = getV(obj)
            v = obj.v;
            t = obj.t;
            U = obj.U; % Stage approximations in previous time step.
            T = obj.T; % Stage times in previous time step.
            K = obj.K; % Stage rates in previous time step.
        end

        function [a,b,c,s] = getTableau(obj)
            a = obj.a;
            b = obj.b;
            c = obj.c;
            s = obj.s;
        end

        % Returns quadrature weights for stages in one time step
        function quadWeights = getTimeStepQuadrature(obj)
            [~, b] = obj.getTableau();
            quadWeights = obj.k*b;
        end

        function obj = step(obj)
            v = obj.v;
            a = obj.a;
            b = obj.b;
            c = obj.c;
            s = obj.s;
            S = obj.S;
            dt = obj.k;
            K = obj.K;
            U = obj.U;
            D = obj.D;
            data = obj.data;

            for i = 1:s
                U(:,i) = v;
                for j = 1:i-1
                    U(:,i) = U(:,i) + dt*a(i,j)*K(:,j);
                end

                K(:,i) = D*U(:,i);
                obj.T(i) = obj.t + c(i)*dt;

                % Data from continuous function and discrete time-points.
                if ~isempty(S)
                    K(:,i) = K(:,i) + S(obj.T(i));
                end
                if ~isempty(data)
                    K(:,i) = K(:,i) + data(:,obj.n*s + i);
                end

            end

            obj.v = v + dt*K*b;
            obj.t = obj.t + dt;
            obj.n = obj.n + 1;
            obj.U = U;
            obj.K = K;
        end
    end


    methods (Static)
        function k = getTimeStep(lambda, order)
            default_arg('order', 4);
            switch order
            case 4
                k = time.rk4.get_rk4_time_step(lambda);
            otherwise
                error('Time-step function not available for this order');
            end
        end
    end

end