classdef Rewards2 < handle
    properties
        stop_sign_d_bndry = 10
        goal_d_bndry = 10
        
        % acceleration quality boundaries:
        % 0---kACC_L----------kACC_LC----------kACC_CH----------kACC_H
        %            light^     comfortable^           hard^
%         kACC_L = .25;
%         kACC_LC = 1.5;
%         kACC_CH = 2;
%         kACC_H = 3;

        kACC_L = 0;%.25;
        kACC_LC = 1.5;
        kACC_CH = 2;
        kACC_H = 3;

        
        kACC_BAD = 0;
        kACC_LIGHT = 1;
        kACC_COMFORTABLE = 2;
        kACC_HARD = 3;
        kDEACC_LIGHT = 4;
        kDEACC_COMFORTABLE = 5;
        kDEACC_HARD = 6;
        kMAINTAIN = 7
        
        % velocity quality boundaries:
        % 0----------kVEL_SP----------kVEL_PF----------kVEL_F
        %  slow^        preferred^           fast^
        kVEL_S = 0;
        kVEL_SP = 4.5;
        kVEL_PF = 5.5;
        kVEL_F = 7;

        kVEL_BAD = 0;
        kVEL_SLOW = 1;
        kVEL_PREFERRED = 2;
        kVEL_FAST = 3;
        kVEL_STOPPED = 4;
        
        % reward multipliers
        kVERYHIGH_MULTIPLIER = 8;
        kHIGH_MULTIPLIER = 4;
        kMEDIUM_MULTIPLIER = 2;
        kLOW_MULTIPLIER = 1;
    
        v_des
    end
    
    methods
        function obj = Rewards2(v_des)
            obj.v_des = v_des;
        end
          
        function reward = compute(obj, v_from, a, params)
            reward = 0;
            temp = [];
            temp2 = [];
            if isfield(params, 'v_to')
                reward = reward + obj.driveDesired(v_from, a, params.v_to);
                temp2 = [temp2, 'driving '];
                temp = [temp, obj.driveDesired(v_from, a, params.v_to)];
            end
            if isfield(params, 'd_goal') && isfield(params, 'v_to')
                reward = reward + obj.goalAhead(v_from, a, params.d_goal, params.v_to);
                temp2 = [temp2, 'goal '];
                temp = [temp, obj.goalAhead(v_from, a, params.d_goal, params.v_to)];
            end
            if isfield(params, 'in_goal')
                if params.in_goal
                    reward = reward + obj.inGoal(v_from, a);
                    temp2 = [temp2, 'in_goal '];
                    temp = [temp, obj.inGoal(v_from, a)];
                end 
            end
            if isfield(params, 'd_stop_sign') && isfield(params, 'v_to')
                reward = reward + obj.stopAhead(v_from, a, params.d_stop_sign, params.v_to);
                temp2 = [temp2, 'stop_sign '];
                temp = [temp, obj.stopAhead(v_from, a, params.d_stop_sign, params.v_to)];
            end    
            if isfield(params, 'in_stop_no_go')
                if params.in_stop
                    reward = reward + obj.inStopNoGo(v_from, a);
                    temp2 = [temp2, 'in_stop '];
                    temp = [temp, obj.inStop(v_from, a)];
                end 
            end    
            if isfield(params, 'in_stop_go')
                if params.in_stop
                    reward = reward + obj.inStopGo(v_from, a);
                    temp2 = [temp2, 'in_stop '];
                    temp = [temp, obj.inStop(v_from, a)];
                end 
            end    
            
                
            
%             fprintf('rewards: %2.3f, %2.3f\n', obj.driveDesired(v_from, a, params.v_to), obj.goalAhead(v_from, a, params.d_goal));
%             disp(temp2)
%             disp(temp)
%             
        end
        
        function reward = driveDesired(obj, v_from, a, v_to)
            % The context multiplier;
            reward = obj.kLOW_MULTIPLIER; 
            % reward function
            reward = reward * (1-abs(obj.v_des - v_to)/obj.v_des)^2;
            
            v_q = obj.velQuality(v_from);
            a_q = obj.accQuality(a);
            if (v_q == obj.kVEL_PREFERRED && a_q == obj.kMAINTAIN)
                reward = reward * obj.kVERYHIGH_MULTIPLIER;
            elseif (v_q == obj.kVEL_PREFERRED && a_q == obj.kACC_LIGHT) ...
                    || (v_q == obj.kVEL_PREFERRED && a_q == obj.kDEACC_LIGHT) ...
                    || (v_q == obj.kVEL_SLOW && a_q == obj.kACC_COMFORTABLE) ...
                    || (v_q == obj.kVEL_STOPPED && a_q == obj.kACC_COMFORTABLE) ...
                    || (v_q == obj.kVEL_FAST && a_q == obj.kDEACC_COMFORTABLE)
                reward = reward * obj.kHIGH_MULTIPLIER;
            elseif (v_q == obj.kVEL_PREFERRED && a_q == obj.kACC_COMFORTABLE) ...
                    || (v_q == obj.kVEL_PREFERRED && a_q == obj.kDEACC_COMFORTABLE)
                reward = reward * obj.kLOW_MULTIPLIER;
            elseif (v_q ~= obj.kVEL_STOPPED || a_q ~= obj.kMAINTAIN)
                reward = reward * obj.kLOW_MULTIPLIER;
            else
                reward = 0;
            end
            
        end
        
        function reward = stopAhead(obj, v_from, a, d_to_stop, v_to)
            d_bndry = obj.stop_sign_d_bndry;
            if d_to_stop > d_bndry
                % too far away to care about 
                reward = 0;
                return
            end
            % The context multiplier
            reward = obj.kHIGH_MULTIPLIER;
            % reward function
            v_decel_des = obj.v_des/d_bndry*d_to_stop;
            reward = reward * (1-abs(v_decel_des - v_to)/obj.v_des)^2;
                       
%             if v_to > v_decel_des
%                 reward = 0;
%             end
            
            v_q = obj.velQuality(v_from);
            a_q = obj.accQuality(a);
            
            if (a_q == obj.kDEACC_COMFORTABLE)% ...
                    %|| (v_q == obj.kVEL_SLOW && a_q == obj.kDEACC_LIGHT)
                reward = reward * obj.kHIGH_MULTIPLIER;
            elseif (a_q == obj.kDEACC_LIGHT)
                reward = reward * obj.kMEDIUM_MULTIPLIER;
            else
                reward = 0;
            end
        end
        
        function reward = goalAhead(obj, v_from, a, d_to_goal, v_to)
            d_bndry = obj.goal_d_bndry;
            if d_to_goal > d_bndry
                % too far away to care about 
                reward = 0;
                return
            end
            % The context multiplier
            reward = obj.kHIGH_MULTIPLIER; 
            % reward function
            v_decel_des = obj.v_des/d_bndry*d_to_goal;
            
            reward = reward * (1-abs(v_decel_des - v_to)/obj.v_des)^2;
            v_q = obj.velQuality(v_from);
            a_q = obj.accQuality(a);
            
            % driving
            if (a_q == obj.kDEACC_COMFORTABLE) %...
                   % || (v_q == obj.kVEL_SLOW && a_q == obj.kDEACC_LIGHT)
                reward = reward * obj.kHIGH_MULTIPLIER;
            elseif (a_q == obj.kDEACC_LIGHT)
                reward = reward * obj.kMEDIUM_MULTIPLIER;
            else
                reward = 0;                
            end
        end
        
        
        function reward = inGoal(obj, v_from, a)
            % The context multiplier
            reward = obj.kHIGH_MULTIPLIER; 
            % No reward function
            
            v_q = obj.velQuality(v_from);
            a_q = obj.accQuality(a);
            
            if (v_q == obj.kVEL_STOPPED && a_q == obj.kMAINTAIN)
                reward = reward * obj.kVERYHIGH_MULTIPLIER;
            elseif (v_q == obj.kVEL_SLOW && a_q == obj.kDEACC_COMFORTABLE)
                reward = reward * obj.kHIGH_MULTIPLIER;
            elseif (v_q == obj.kVEL_SLOW && a_q == obj.kDEACC_LIGHT)
                reward = reward * obj.kMEDIUM_MULTIPLIER;
            else
                reward = 0;
%             elseif (v_q == obj.kVEL_SLOW && a_q == obj.kDEACC_HARD)
%                 reward = reward * obj.kLOW_MULTIPLIER;
            end
        end

        function reward = inStopNoGo(obj, v_from, a)
            % The context multiplier
            reward = obj.kHIGH_MULTIPLIER; 
            % No reward function
            
            v_q = obj.velQuality(v_from);
            a_q = obj.accQuality(a);

            if (v_q == obj.kVEL_STOPPED && a_q == obj.kMAINTAIN)
                reward = reward * obj.kVERYHIGH_MULTIPLIER;
            elseif (v_q == obj.kVEL_SLOW && a_q == obj.kDEACC_COMFORTABLE)...
                    || (v_q == obj.kVEL_SLOW && a_q == obj.kDEACC_LIGHT)
                reward = reward * obj.kHIGH_MULTIPLIER;
            else
                reward = 0;
%             elseif (v_q == obj.kVEL_SLOW && a_q == obj.kDEACC_LIGHT)
%                 reward = reward * obj.kLOW_MULTIPLIER;
%             elseif (v_q == obj.kVEL_SLOW && a_q == obj.kDEACC_HARD)
%                 reward = reward * obj.kLOW_MULTIPLIER;
            end
        end
        
        function reward = inStopGo(obj, v_from, a)
            % The context multiplier
            reward = obj.kHIGH_MULTIPLIER; 
            % No reward function
            
            v_q = obj.velQuality(v_from);
            a_q = obj.accQuality(a);

            if (v_q == obj.kVEL_SLOW && a_q == obj.kACC_COMFORTABLE)
                reward = reward * obj.kVERYHIGH_MULTIPLIER;
            elseif (v_q == obj.kVEL_SLOW && a_q == obj.kACC_LIGHT)
                reward = reward * obj.kMEDIUM_MULTIPLIER;
            else
                reward = 0;
%             elseif (v_q == obj.kVEL_SLOW && a_q == obj.kDEACC_LIGHT)
%                 reward = reward * obj.kLOW_MULTIPLIER;
%             elseif (v_q == obj.kVEL_SLOW && a_q == obj.kDEACC_HARD)
%                 reward = reward * obj.kLOW_MULTIPLIER;
            end
        end
        
        
        function acc_quality = accQuality(obj, a)
            if a == 0
                acc_quality = obj.kMAINTAIN;
            elseif -obj.kACC_H < a && a < -obj.kACC_CH
                acc_quality = obj.kDEACC_HARD;
            elseif -obj.kACC_CH < a && a < -obj.kACC_LC
                acc_quality = obj.kDEACC_COMFORTABLE;
            elseif -obj.kACC_LC < a && a < -obj.kACC_L
                acc_quality = obj.kDEACC_LIGHT;
            elseif obj.kACC_L < a && a < obj.kACC_LC
                acc_quality = obj.kACC_LIGHT;
            elseif obj.kACC_LC < a && a < obj.kACC_CH
                acc_quality = obj.kACC_COMFORTABLE;
            elseif obj.kACC_CH < a && a < obj.kACC_H
                acc_quality = obj.kACC_HARD;
            else
                acc_quality = obj.kACC_BAD;
            end
        end
        
        function vel_quality = velQuality(obj, v)
            if v == 0
                vel_quality = obj.kVEL_STOPPED;
            elseif obj.kVEL_S < v && v < obj.kVEL_SP
                vel_quality = obj.kVEL_SLOW;
            elseif obj.kVEL_SP < v && v < obj.kVEL_PF
                vel_quality = obj.kVEL_PREFERRED;
            elseif obj.kVEL_PF < v && v < obj.kVEL_F
                vel_quality = obj.kVEL_FAST;
            else
                vel_quality = obj.kVEL_BAD;
            end
        end
        
    end
    
end

