classdef (StrictDefaults) MultipleStageBlockSystem < matlab.System
    % This function is unsupported and might change or be removed without
    % notice in a future version.

    %#codegen

    properties (Nontunable)
        blkData     = struct('p',1,'nmd',1,'nx',1,'nmv',1,'npara',1,'nP',1,'nZ',1,'Ts',1)
        TsControl   = -1
        mexname     = ''

        UseObjectTs             (1,1) logical = false
        UseMEX                  (1,1) logical = false
        md_enabled              (1,1) logical = false
        stateparam_enabled      (1,1) logical = false
        stageparam_enabled      (1,1) logical = false
        cost_enabled            (1,1) logical = false
        mvseq_enabled           (1,1) logical = false
        stateseq_enabled        (1,1) logical = false
        slack_enabled           (1,1) logical = false
        status_enabled          (1,1) logical = false
        mv_min                  (1,1) logical = false
        mv_max                  (1,1) logical = false
        mvrate_min              (1,1) logical = false
        mvrate_max              (1,1) logical = false
        state_min               (1,1) logical = false
        state_max               (1,1) logical = false
        terminal_state          (1,1) logical = false
        nlp_initialize          (1,1) logical = false
    end

    properties (Nontunable, Access = private)
        blkDataInternal
    end

    % Pre-computed constants or internal states
    properties (Access = private)
        onlinedata
    end

    properties (Hidden, Access=private)
        isInitialized = false
    end

    methods
        % Constructor
        function obj = MultipleStageBlockSystem(varargin)
            setProperties(obj,nargin,varargin{:})
        end
    end

    methods (Access = protected)
        function setupImpl(obj,x,lastMV,~,~,~,~,~,~,~,~,~,~,~)
            % Perform one-time calculations, such as computing constants
            signals = MultipleStageBlockSystem.getSimulationData(obj, obj.blkData.nZ, x);
            obj.blkDataInternal = coder.const(mpc.internal.cgmres.setDimensions(x, lastMV, obj.blkData, signals));
            obj.onlinedata = MultipleStageBlockSystem.getSimulationData(obj, obj.blkDataInternal.nZ, x);
        end

        function [mv, Cost, MVopt, Xopt, Slack, ExitFlag] = stepImpl(obj, x, lastMV, md, MVMin, MVMax, MVRateMin, MVRateMax, StateMin, StateMax, StateParameter, StageParameter, TerminalState, extZ0)
            %STEPIMPL Implementation of the step method for the MPC controller
            %
            %   [__] = stepImpl(obj,__) takes the current state X, the last
            %   manipulated variables LASTMV, measured disturbances MD, and
            %   various constraints and parameters to compute the optimal
            %   MV for the current step, along with other related
            %   information such as the cost, optimal MV and state
            %   sequences, slack variables, and an exit flag from the
            %   optimization solver.
            %
            %   Inputs:
            %       obj             - The MPC controller object
            %       x               - The current state vector
            %       lastMV          - The last manipulated variables used
            %       md              - Measured disturbances
            %       MVMin           - Lower bounds for MV
            %       MVMax           - Upper bounds for MV
            %       MVRateMin       - Lower bounds for the rate of change of MV
            %       MVRateMax       - Upper bounds for the rate of change of MV
            %       StateMin        - Lower bounds for states
            %       StateMax        - Upper bounds for states
            %       StateParameter  - Parameters associated with states
            %       StageParameter  - Parameters associated with each stage of optimization
            %       TerminalState   - Terminal state for the prediction horizon
            %       extZ0           - External initial guess for the optimizer
            %
            %   Outputs:
            %       mv              - Optimal manipulated variables for the current step
            %       Cost            - Cost associated with the optimal solution
            %       MVopt           - Optimal sequence of MV over the prediction horizon
            %       Xopt            - Optimal sequence of states over the prediction horizon
            %       Slack           - Slack variables associated with the optimization
            %       ExitFlag        - Exit flag from the optimization solver

            % Update simulation data with the given parameters
            obj.updateSimulationData(md, MVMin, MVMax, MVRateMin, MVRateMax, StateMin, StateMax, StateParameter, StageParameter, TerminalState, extZ0);
            
            % Initialization for the first simulation step
            if ~obj.isInitialized && ~obj.nlp_initialize
                obj.isInitialized = true;
                
                % Solver-specific initialization
                if strcmpi(obj.blkDataInternal.solver, 'fmincon')
                    % Initialize with a guess based on the x and lastMV
                    if obj.blkDataInternal.hasmvrate
                        obj.onlinedata.InitialGuess =  mpc.internal.nlmpcMultistage.getInitialGuessXMVDMVE(obj.blkDataInternal,[],x,lastMV);
                    else
                        obj.onlinedata.InitialGuess =  mpc.internal.nlmpcMultistage.getInitialGuessXMVE(obj.blkDataInternal,[],x,lastMV);
                    end                    
                elseif strcmpi(obj.blkDataInternal.solver, 'cgmres')
                    % Cold start initialization for CGMRES solver
                    [Z, dZdt] = mpc.internal.cgmres.coldStartInitialization(x, lastMV, obj.blkDataInternal, obj.onlinedata);
                    obj.onlinedata.InitialGuess = [Z; dZdt; zeros(obj.blkDataInternal.nx,1)];
                end
            end

            % Execute the optimization to compute the optimal MV
            if coder.target('MATLAB')
                if obj.UseMEX
                    % Use MEX function if specified for faster execution
                    mexfcn = str2func(obj.blkDataInternal.mexname);
                    try
                        [mv, obj.onlinedata, info] = mexfcn(x, lastMV, obj.onlinedata);
                    catch ME
                        throw(ME)
                    end
                else
                    % Use MATLAB nlmpcmove function if not using MEX
                    msobj = slResolve(obj.blkDataInternal.nlmpcobj,obj.blkDataInternal.blk);
                    try
                        [mv, obj.onlinedata, info] = nlmpcmove(msobj, x, lastMV, obj.onlinedata);
                    catch ME
                        throw(ME)
                    end
                end
            else 
                % Code generation path for deploying the controller outside of MATLAB
                [mv, obj.onlinedata, info] = nlmpcmoveCodeGeneration(obj.blkDataInternal, x, lastMV, obj.onlinedata);
            end

            % Extract information from the optimization output
            Cost     = info.Cost;
            MVopt    = info.MVopt;
            Xopt     = info.Xopt;
            Slack    = zeros(obj.blkData.p+1, obj.blkData.nx);
            ExitFlag = info.ExitFlag;
        end

        function resetImpl(obj)
            % Initialize / reset internal properties
        end

        function validateInputsImpl(obj,x,lastMV,~,~,~,~,~,~,~,~,~,~,extZ0)
            %VALIDATEINPUTSIMPL Validate inputs to the step method at initialization
            %   This function is responsible for validating the external
            %   initial guess EXTZ0 provided to the step method of the MPC
            %   controller when the 'cgmres' solver is used.

            % Check if the 'cgmres' solver is used and the controller is initializing
            if isfield(obj.blkData, 'solver') && strcmpi(obj.blkData.solver, 'cgmres') && obj.nlp_initialize
                if isempty(obj.blkDataInternal)
                    % Internal block data is not yet initialized
                    simdata  = MultipleStageBlockSystem.getSimulationData(obj, obj.blkData.nZ, x);
                    coredata = coder.const(mpc.internal.cgmres.setDimensions(x, lastMV, obj.blkData, simdata));
                    validateattributes(extZ0,{'double'},{'vector','real','finite','numel',coredata.nZ},'nlmpcmove','extZ0');
                else
                    % Internal block data is already initialized
                    validateattributes(extZ0,{'double'},{'vector','real','finite','numel',obj.blkDataInternal.nZ},'nlmpcmove','extZ0');
                end
            end
        end

        function flag = isInputSizeMutableImpl(obj,index)
            % Set input size cannot change between calls to the System object
            flag = false;
        end

        function [out1,out2,out3,out4,out5,out6] = getOutputSizeImpl(obj)
            % Return size for each output port
            out1 = propagatedInputSize(obj,2);
            out2 = [1 1];
            out3 = [obj.blkData.p+1, obj.blkData.nmv];
            out4 = [obj.blkData.p+1, obj.blkData.nx];
            out5 = [obj.blkData.p+1, obj.blkData.nx];
            out6 = [1 1];
        end

        function [out1,out2,out3,out4,out5,out6] = getOutputDataTypeImpl(obj)
            % Return data type for each output port
            out1 = "double";
            out2 = "double";
            out3 = "double";
            out4 = "double";
            out5 = "double";
            out6 = "double";
        end

        function [out1,out2,out3,out4,out5,out6] = isOutputComplexImpl(obj)
            % Return true for each output port with complex data
            out1 = false;
            out2 = false;
            out3 = false;
            out4 = false;
            out5 = false;
            out6 = false;
        end

        function [out1,out2,out3,out4,out5,out6] = isOutputFixedSizeImpl(obj)
            % Return true for each output port with fixed size
            out1 = true;
            out2 = true;
            out3 = true;
            out4 = true;
            out5 = true;
            out6 = true;
        end

        function sts = getSampleTimeImpl(obj)
            %GETSAMPLETIMEIMPL Define sample time for the System object
            %   This function specifies the sample time used by the MPC
            %   controller System object. The sample time can be set to a
            %   user-defined value, inherited from the upstream block, or
            %   matched to the internal sample time of the object.
            
            if obj.UseObjectTs
                sts = obj.createSampleTime("Type", "Discrete", "SampleTime", obj.blkData.Ts);
            elseif obj.TsControl ~= -1
                sts = obj.createSampleTime("Type", "Discrete", "SampleTime", obj.TsControl);
            else
                sts = obj.createSampleTime('Type','Inherited');
            end
        end
    end

    methods (Access = protected)

        function updateSimulationData(obj, md, MVMin, MVMax, MVRateMin, MVRateMax, StateMin, StateMax, StateParameter, StageParameter, TerminalState, extZ0)

            p = obj.blkDataInternal.p;
            if obj.md_enabled
                obj.onlinedata.MeasuredDisturbance = MultipleStageBlockSystem.extendBounds(md,p+1,obj.blkDataInternal.nmd);
            end
            if obj.mv_min
                obj.onlinedata.MVMin = MultipleStageBlockSystem.extendBounds(MVMin,p,obj.blkDataInternal.nmv);
            end
            if obj.mv_max
                obj.onlinedata.MVMax = MultipleStageBlockSystem.extendBounds(MVMax,p,obj.blkDataInternal.nmv);
            end
            if obj.mvrate_min
                obj.onlinedata.MVRateMin = MultipleStageBlockSystem.extendBounds(MVRateMin,p,obj.blkDataInternal.nmv);
            end
            if obj.mvrate_max
                obj.onlinedata.MVRateMax = MultipleStageBlockSystem.extendBounds(MVRateMax,p,obj.blkDataInternal.nmv);
            end
            if obj.state_min
                obj.onlinedata.StateMin = MultipleStageBlockSystem.extendBounds(StateMin,p,obj.blkDataInternal.nx);
            end
            if obj.state_max
                obj.onlinedata.StateMax = MultipleStageBlockSystem.extendBounds(StateMax,p,obj.blkDataInternal.nx);
            end
            if obj.stateparam_enabled
                obj.onlinedata.StateFcnParameter = StateParameter;
            end
            if obj.stageparam_enabled
                obj.onlinedata.StageParameter = StageParameter;
            end
            if obj.terminal_state
                obj.onlinedata.TerminalState = TerminalState;
            end
            if obj.nlp_initialize
                obj.onlinedata.InitialGuess = extZ0;
            end
        end
    end

    methods (Static)
        function B = extendBounds(A,maxrows,n)
            arguments
                A
                maxrows
                n = []
            end

            if iscolumn(A) && numel(A) == n
                A = A.';
            end
            rows = size(A,1);
            if rows < maxrows
                B = [A; repmat(A(end, :), maxrows-rows, 1)];
            else 
                B = A(1:maxrows,:); 
            end
        end

        function onlinedata = getSimulationData(maskObject, nZ, x)
            %GETSIMULATIONDATA generates default "simdata" structure to be used with
            %"nlmpcmove" for simulation.

            %#codegen

            coder.allowpcode('plain');
            coder.internal.prefer_const(maskObject, nZ, x);

            onlinedata = struct();
            if maskObject.md_enabled
                onlinedata.MeasuredDisturbance = zeros(1,maskObject.blkData.nmd);
            end
            if maskObject.stateparam_enabled
                onlinedata.StateFcnParameter = zeros(maskObject.blkData.npara,1);
            end
            if maskObject.stageparam_enabled
                onlinedata.StageParameter = zeros(maskObject.blkData.nP,1);
            end
            if maskObject.slack_enabled
                onlinedata.StageParameter = zeros(maskObject.blkData.nE,1);
            end
            if maskObject.mv_min
                onlinedata.MVMin = zeros(maskObject.blkData.p,maskObject.blkData.nmv);
            end
            if maskObject.mv_max
                onlinedata.MVMax = zeros(maskObject.blkData.p,maskObject.blkData.nmv);
            end
            if maskObject.mvrate_min
                onlinedata.MVRateMin = zeros(maskObject.blkData.p,maskObject.blkData.nmv);
            end
            if maskObject.mvrate_max
                onlinedata.MVRateMax = zeros(maskObject.blkData.p,maskObject.blkData.nmv);
            end
            if maskObject.state_min
                onlinedata.StateMin = zeros(maskObject.blkData.p,maskObject.blkData.nx);
            end
            if maskObject.state_max
                onlinedata.StateMax = zeros(maskObject.blkData.p,maskObject.blkData.nx);
            end
            if maskObject.terminal_state
                onlinedata.TerminalState = zeros(maskObject.blkData.nx, 1);
            end
            % if obj.nlp_initialize
            %     onlinedata.InitialGuess = zeros(nZ, 1);
            % end
            onlinedata.InitialGuess = zeros(nZ, 1);
        end

    end

    methods(Access = protected,Static)
        function flag = showSimulateUsingImpl
            flag = false;
        end
        function simMode = getSimulateUsingImpl()
            simMode = "Interpreted execution";
            % simMode = "Code generation";
        end
    end

end