classdef GoalRewards < handle
    properties
        goal_d_bndry = 10

        r
        scene
    end
    
    methods
        function obj = GoalRewards(scene, reward_stats, ~, ~)
            obj.scene = scene;
            obj.r = reward_stats;
            
        end
          
        function reward = compute(obj, t, x, u, x_to, ~) % v_from, a, params)
            d_goal = obj.scene.distToGoal([x_to(1), 0]);
            d_goal_center = obj.scene.distToGoalCenter([x_to(1), 0]);
            in_goal = false;
            if d_goal < 0.0
                in_goal = true;
            end

            reward = 0;
            if ~in_goal
                reward = reward + obj.goalAhead(x(4), u, d_goal_center, x_to(4));
            else % in_goal
                reward = reward + obj.inGoal(x(4), u);
            end    
        end
        
        
        function reward = goalAhead(obj, v_from, a, d_to_goal_center, v_to)
            d_bndry = obj.goal_d_bndry;
            if d_to_goal_center > d_bndry
                % too far away to care about 
                reward = 0;
                return
            end
            % The context multiplier
            reward = obj.r.kHIGH_MULTIPLIER; 
            % reward function
            v_decel_des = obj.r.v_des/d_bndry*d_to_goal_center;
            
%             if v_to > v_decel_des
%                 reward = 0;
%                 return 
%             end
            
            reward = reward * (1-abs(v_decel_des - v_to)/obj.r.v_des)^2;
            
            v_q = obj.r.velQuality(v_from);
            a_q = obj.r.accQuality(a);
            
            
            
            % driving
            if (a_q == obj.r.kDEACC_COMFORTABLE) %...
                   % || (v_q == obj.kVEL_SLOW && a_q == obj.kDEACC_LIGHT)
                reward = reward * obj.r.kHIGH_MULTIPLIER;
            elseif (a_q == obj.r.kDEACC_LIGHT)
                reward = reward * obj.r.kHIGH_MULTIPLIER;
            else
                reward = 0;
            end
        end
        
        
        function reward = inGoal(obj, v_from, a)
            % The context multiplier
            reward = obj.r.kHIGH_MULTIPLIER; 
            % No reward function
            
            v_q = obj.r.velQuality(v_from);
            a_q = obj.r.accQuality(a);
            
            if (v_q == obj.r.kVEL_STOPPED && a_q == obj.r.kMAINTAIN)
                reward = reward * obj.r.kVERYHIGH_MULTIPLIER;
            elseif (v_q == obj.r.kVEL_SLOW && a_q == obj.r.kDEACC_COMFORTABLE)...
                    || (v_q == obj.r.kVEL_SLOW && a_q == obj.r.kDEACC_LIGHT)
                reward = reward * obj.r.kHIGH_MULTIPLIER;
            elseif (v_q == obj.r.kVEL_SLOW && a_q == obj.r.kDEACC_HARD)
                reward = reward * obj.r.kLOW_MULTIPLIER;
            else
                reward = 0;
            end
            
            if (v_q == obj.r.kVEL_STOPPED && a_q == obj.r.kMAINTAIN)
                reward = reward * obj.r.kVERYHIGH_MULTIPLIER;
            elseif (v_q == obj.r.kVEL_SLOW && a_q == obj.r.kDEACC_COMFORTABLE)
                reward = reward * obj.r.kHIGH_MULTIPLIER;
            elseif (v_q == obj.r.kVEL_SLOW && a_q == obj.r.kDEACC_LIGHT)
                reward = reward * obj.r.kMEDIUM_MULTIPLIER;
            elseif (v_q == obj.r.kVEL_SLOW && a_q == obj.r.kDEACC_HARD)
                reward = reward * obj.r.kLOW_MULTIPLIER;
            else                
                reward = 0;
            end
        end
    end
    
end

