classdef General < time.Timestepper
    properties
        F       % RHS of the ODE
        dt      % Time step
        t       % Time point
        v       % Solution vector
        n       % Time level
        scheme  % The scheme used for the time stepping, e.g rk4, rk6 etc.
        bt
        V       % All stage approximations in most recent time step
        K       % All stage rates in most recent time step
    end


    methods
        % Timesteps v_t = F(t,v), using the specified ButcherTableau
        % from t = t0 with timestep dt and initial conditions v(0) = v0
        function obj = General(F, dt, t0, v0, bt)
            assertType(bt, 'time.rk.ButcherTableau')
            obj.F = F;
            obj.dt = dt;
            obj.t = t0;
            obj.v = v0;
            obj.n = 0;

            assert(bt.isExplicit())
            obj.bt = bt;
        end

        % v: Current solution
        % t: Current time
        % V: All stage approximations in most recent time step
        % K: All stage rates in most recent time step
        % T: Time points (corresponding to V and K) in most recent time step
        function [v,t] = getV(obj)
            v = obj.v;
            t = obj.t;
        end

        function obj = step(obj)
            s = obj.bt.nStages();
            a = obj.bt.a;
            b = obj.bt.b;
            c = obj.bt.c;

            % Compute rates K
            K = zeros(length(v), s);
            for i = 1:s
                V_i = obj.v;
                for j = 1:i-1
                    V_i = V_i + dt*a(i,j)*K(:,j);
                end
                K(:,i) = F(t+dt*c(i), V_i);
            end

            % Compute updated solution
            v_next = v;
            for i = 1:s
                v_next = v_next + dt*b(i)*K(:,i);
            end

            obj.v = v_next;
            obj.t = obj.t + obj.dt;
            obj.n = obj.n + 1;
        end

        % TBD: Method name
        % TBD: Parameter name
        %
        % Takes a regular step but with discreteRates(:,i) added to RHS for stage i.
        %  v_t = F(t,v) + discreteRates(:, ...)
        %
        % Also returns the stage approximations (V) and stage rates (K).
        function [v,t, V, K] = stepWithDiscreteData(obj, discreteRates)
            s = obj.bt.nStages();
            a = obj.bt.a;
            b = obj.bt.b;
            c = obj.bt.c;

            % Compute rates K and stage approximations V
            K = zeros(length(v), s);
            V = zeros(length(v), s);
            for i = 1:s
                V_i = obj.v;
                for j = 1:i-1
                    V_i = V_i + dt*a(i,j)*K(:,j);
                end

                K_i = F(t+dt*c(i), V_i);
                K_i = K_i + discreteRates(:,i);

                V(:,i) = V_i;
                K(:,i) = K_i;
            end

            % Compute updated updated solution
            v_next = v;
            for i = 1:s
                v_next = v_next + dt*b(i)*K(:,i);
            end

            obj.v = v_next;
            obj.t = obj.t + obj.dt;
            obj.n = obj.n + 1;
        end

        % Returns a vector of time points, including substage points,
        % in the time interval [t0, tEnd].
        % The time-step obj.dt is assumed to be aligned with [t0, tEnd] already.
        function tvec = timePoints(obj, t0, tEnd)
            % TBD: Should this be implemented here or somewhere else?
            N = round( (tEnd-t0)/obj.dt );
            tvec = zeros(N*obj.s, 1);
            s = obj.coeffs.s;
            c = obj.coeffs.c;
            for i = 1:N
                ind = (i-1)*s+1 : i*s;
                tvec(ind) = ((i-1) + c')*obj.dt;
            end
        end

        % Returns a vector of quadrature weights corresponding to grid points
        % in time interval [t0, tEnd], substage points included.
        % The time-step obj.dt is assumed to be aligned with [t0, tEnd] already.
        function weights = quadWeights(obj, t0, tEnd)
            % TBD: Should this be implemented here or somewhere else?
            N = round( (tEnd-t0)/obj.dt );
            b = obj.coeffs.b;
            weights = repmat(b', N, 1);
        end
    end

    methods(Static)
        % TBD: Function name
        function ts = methodFromStr(F, dt, t0, v0, methodStr, discreteData)
            default_arg('discreteData', []);

            try
                bt = time.rk.ButcherTableau.(method);
            catch
                error('Runge-Kutta method ''%s'' is not implemented', methodStr)
            end

            ts = time.rk.General(F, dt, t0, v0, bt, discreteData);
        end
    end
end