classdef ConstraintFunctions < handle
% Checks constraints on the state and control trajectories as well as 
% symbolic task constraints handled through an automata generated from an 
% LTL formula and predicates (propositions).
%
% Initialize with the inequality constraints given by cs.parameter where 
% parameter is associated with
%   xBound: NS x 2 matrix of upper and lower bounds to the state.
%   uBound: NI x 2 matrix of upper and lower bounds to the control input.
%   maintainRoad: true or false 
% Notes: 
%   use (-)inf if unbounded (lower or) upper bound. 
%   neglect a constraint parameter to not check that constraint.
%
% Call compute(ts, xs, us) with state and control trajectories xs and us
% for time ts.
    properties
        scene
        cs
        proposition_handlers
        LTL_automatas
        check_LTL = false
        
        internal_proposition_state_out 
    end
    
    methods
        function obj = ConstraintFunctions(scenario, constraints, LTL_automatas, proposition_handlers)
        % constraints are a struct of xBounds, uBounds and whether to 
        %   maintainRoad or not. 
        % LTL_automata is the object for stepping through the LTL specified
        %   automata for symbolic task constraints.
        % proposition_handler transforms a trajectory into a sequence of
        %   booleans for the LTL_automata to check.
            obj.scene = scenario;
            obj.cs = constraints;
            if nargin>2
                assert(nargin>3)
                obj.check_LTL = true;
                obj.LTL_automatas = LTL_automatas;
                obj.proposition_handlers = proposition_handlers;
            end
        end
        
        function is_satisfied = compute(obj, ts, xs, us, symbolic_state, internal_proposition_state)
        % Checks that the constraints specified in cs are satisfied    
            if isfield(obj.cs,'xBounds')
                if ~obj.satisfyxBounds(xs)
                    is_satisfied = false;
                    return
                end
            end
            if isfield(obj.cs,'uBounds')
                if ~obj.satisfyuBounds(us)
                    is_satisfied = false;
                    return
                end
            end
            if isfield(obj.cs,'maintainRoad')
                if obj.cs.maintainRoad % could be set to false.
                    if ~obj.stayOnRoad(xs)
                        is_satisfied = false;
                        return
                    end
                end
            end
            if obj.check_LTL
                if nargin >= 6
                    satisfiesLTL = obj.satisfyLTL(ts, xs, us, symbolic_state, internal_proposition_state);
                else
                    satisfiesLTL = obj.satisfyLTL(ts, xs, us, symbolic_state);
                end
                if ~satisfiesLTL
                    is_satisfied = false;
                    return
                end
            end
            is_satisfied = true;
        end
        
        function is_satisfied = satisfyLTL(obj, ts, xs, us, symbolic_state, internal_proposition_state)
            % For some reason this is illegal:
            % @obj.proposition_handler.traj2Propostions
            for i = 1:numel(obj.LTL_automatas)
                prop_handler = obj.proposition_handlers{i};
                if nargin >= 6
                    is_satisfied = obj.LTL_automatas{i}.check_run(ts, xs, us, @prop_handler.traj2Propositions, symbolic_state{i}, internal_proposition_state{i});
                    obj.internal_proposition_state_out{i} = obj.LTL_automatas{i}.internal_proposition_state_out;
                else
                    is_satisfied = obj.LTL_automatas{i}.check_run(ts, xs, us, @prop_handler.traj2Propositions, symbolic_state{i});
                end
                if is_satisfied == false
                    return
                end
            end
        end
        
        function is_satisfied = satisfyxBounds(obj, xs)
        % Checks that state trajectory stays inside xBounds.
        
            for i = 1:size(obj.cs.xBounds,1)
                if obj.cs.xBounds(i,1) > -inf
                    if min(xs(:,i)) < obj.cs.xBounds(i,1)
                        is_satisfied = false;
%                         disp(min(xs(:,i)))
%                         disp('x bad')
                        return
                    end
                end
                if obj.cs.xBounds(i,2) < inf
                    if max(xs(:,i)) > obj.cs.xBounds(i,2)
                        is_satisfied = false;
%                         disp(max(xs(:,i)))
%                         disp('x bad')
                        return
                    end
                end
            end
            is_satisfied = true;
        end
        
        function is_satisfied = satisfyuBounds(obj, us)
        % Checks that control input trajectory stays inside uBounds.
            for i = 1:size(obj.cs.uBounds,1)
                if obj.cs.uBounds(i,1) > -inf
                    if min(us(:,i)) < obj.cs.uBounds(i,1)
                        is_satisfied = false;
%                         disp(min(us(:,i)))
%                         disp('u bad')
                        return
                    end
                end
                if obj.cs.uBounds(i,2) < inf
                    if max(us(:,i)) > obj.cs.uBounds(i,2)
                        is_satisfied = false;
%                         disp(max(us(:,i)))
%                         disp('u bad')
                        return
                    end
                end
            end
            is_satisfied = true;
        end
        
        function is_on = stayOnRoad(obj, xs)
        % Checks if ZOOX remains in the road. Only checks the states in
        % obj.xs. No guarantee zoox exits the road between adjacent
        % states in obj.xs.
            
            % check backwards since trajectory end is most likely to be out of
            % bounds.
            for i = size(xs,1):-1:1
                if ~obj.scene.isCarInRoad(xs(i,:))
                    is_on = false;
%                     disp('offroad')
                    return
                end
            end
            is_on = true;
        end
        
    end
    
end