classdef RungekuttaRvMultiStage < time.Timestepper
    properties
        F           % RHS of the ODE
        F_unstable  % RHS of the unstabilized ODE
        k       % Time step
        t       % Time point
        v       % Solution vector
        n       % Time level
        rkScheme  % The particular RK scheme used for time integration
        RV              % Residual Viscosity operator
        DvDt            % Function for computing the time deriative used for the RV evaluation
        v_unstable
        viscosity
    end
    methods

        function obj = RungekuttaRvMultiStage(F, F_unstable, k, t0, v0, RV, DvDt, order)
            obj.F = F;
            obj.F_unstable = F_unstable;
            obj.k = k;
            obj.t = t0;
            obj.v = v0;
            obj.n = 0;
            
            if (order == 4) % Use specialized RK4 scheme
                obj.rkScheme = @time.rk.rungekutta_4;
            else
                % Extract the coefficients for the specified order
                % used for the RK updates from the Butcher tableua.
                [s,a,b,c] = time.rk.butcherTableau(order);
                coeffs = struct('s',s,'a',a,'b',b,'c',c);
                obj.rkScheme = @(v,t,dt,F) time.rk.rungekutta(v, t , dt, F, coeffs);
            end
        
            obj.RV = RV;
            obj.DvDt = DvDt;
            obj.v_unstable = 0*v0;
            obj.viscosity = 0*v0;
        end

        function [v, t] = getV(obj)
            v = obj.v;
            t = obj.t;
        end

        function state = getState(obj)
            dvdt = obj.DvDt(obj.v_unstable);
            [viscosity, Df, firstOrderViscosity, residualViscosity] = obj.RV.evaluate(obj.v, dvdt);
            state = struct('v', obj.v, 'dvdt', dvdt, 'Df', Df, 'viscosity', obj.viscosity, 'residualViscosity', residualViscosity, 'firstOrderViscosity', firstOrderViscosity, 't', obj.t);
        end

        % Advances the solution vector one time step using the Runge-Kutta method given by
        % obj.coeffs, using a fixed residual viscosity for the Runge-Kutta substeps
        function obj = step(obj)            
            
            % Advance solution by unstabilized scheme
            obj.v_unstable = obj.rkScheme(obj.v, obj.t, obj.k, obj.F_unstable);
            % Compute viscosity for current time level based on unstable solution (from next time level)
            % and the current solution.
            obj.viscosity = obj.RV.evaluateViscosity(obj.v, obj.DvDt(obj.v_unstable));
            % Fix the viscosity of the stabilized RHS
            m = length(obj.viscosity);
            F_stable = @(v,t) obj.F(v,t,spdiags(obj.viscosity,0,m,m));
            % Advance solutiont to next time level by stabilized scheme.
            obj.v = obj.rkScheme(obj.v, obj.t, obj.k, F_stable);
            obj.t = obj.t + obj.k;
            obj.n = obj.n + 1;
        end
    end
end