classdef CostFunctions < handle
% Computes the cost for the cost function defined by the given weights ws
% and the state, control trajectories given.

% Initialize with the Weights given by ws.parameter where parameter is
% associated with the metric:
%   centerline: distance from centerline over trajectory.
%   velocity: velocity from zero.
%   acceleration: acceleration from zero
%   steering: steering angle from zero
%   tf: final time
%   XYgoal: distance of final vehicle location from goal point.
%   velocityf: velocity from zero at final time.

% Call compute(ts, xs, us) with state and control trajectories xs and us
% for time ts.
    properties
        scene
        ws
    end
    
    methods
        function obj = CostFunctions(scenario, ws)
            obj.scene = scenario;
            obj.ws = ws;
            if ~isfield(ws, 'centerline')
                obj.ws.centerline = 0;
            end
            if ~isfield(ws, 'XYgoal')
                obj.ws.XYgoal = 0;
            end
            if ~isfield(ws, 'velocityf')
                obj.ws.velocityf = 0;
            end
            if ~isfield(ws, 'velocity')
                obj.ws.velocity = 0;
            end
            if ~isfield(ws, 'velocity_desired')
                % if not set track zero velocity.
                obj.ws.velocity_desired = 0;
            end

            if ~isfield(ws, 'tf')
                obj.ws.tf = 0;
            end
            if ~isfield(ws, 'acceleration')
                obj.ws.acceleration = 0;
            end
            if ~isfield(ws, 'steering')
                obj.ws.steering = 0;
            end
            if ~isfield(ws, 'stop_signs')
                obj.ws.stop_signs = 0;
            end
        end
        
        function cost = compute(obj, ts, xs, us)
        % Computes the cost associated with trajectory given by state xs 
        % control us over times ts.
            cost = 0;
            if obj.ws.centerline ~= 0 || obj.ws.velocity ~= 0
                cost = cost + obj.costLx(ts, xs);
            end
            if obj.ws.acceleration ~= 0 || obj.ws.steering ~= 0
                cost = cost + obj.costLu(ts, us);
            end
            if obj.ws.XYgoal ~= 0 || obj.ws.velocityf ~= 0
                cost = cost + obj.costm(xs(end,:)); 
            end
            if obj.ws.tf ~= 0
                cost = cost + obj.costtf(ts(end));
            end
            if obj.ws.stop_signs ~= 0
                cost = cost + obj.coststopsigns(xs);
            end
        end
        
        function cost = costLx(obj, ts, xs)
        % Computes running cost on distance to centerline and velocity.
            dt =  ts(2:end) - ts(1:end-1);
            cost = 0;
            if obj.ws.centerline ~= 0
                for i = 2:size(xs,1)
                    cost = cost + dt(i-1)*obj.ws.centerline*obj.scene.distToCenterline(xs(i,1:2))^2/2;
%                     cost = cost + dt(i-1)*obj.ws.centerline*expFunc(obj.scene.distToCenterline(xs(i,1:2)))/2;
                end
            end
            if obj.ws.velocity ~= 0
                for i = 2:size(xs,1)
                    verr = xs(i,4) - obj.ws.velocity_desired;
                    cost = cost + dt(i-1)*obj.ws.velocity*verr^2/2;
                end
            end
        end
        function out = expFunc(~, val)
            alpha = 5;
            out = exp(alpha*val) + exp(-alpha*val) - 2;
        end
        function cost = costLu(obj, ts, us)
        % Computes running cost on acceleration and steering angle.
            dt =  ts(2:end) - ts(1:end-1);
            cost = 0;
            if obj.ws.acceleration ~= 0 
                for i = 2:size(us,1)
                    cost = cost + dt(i-1)*obj.ws.acceleration*us(i,1)^2/2;
                end
            end
            if obj.ws.steering ~= 0
                for i = 2:size(us,1)
                    cost = cost + dt(i-1)*obj.ws.steering*us(i,2)^2/2;
                end
            end
        end
        function cost = costm(obj, xf)
        % Computes terminal cost on distance from goal and final velocity.
            err_XY = obj.scene.goal_pt-xf(1:2);
            cost = obj.ws.XYgoal*(dot(err_XY, err_XY))/2;
            cost = cost + obj.ws.velocityf*xf(4)^2/2;
        end
        function cost = costtf(obj, tf)
        % Computes cost on time length.
            cost = obj.ws.tf*tf^2/2;
        end
        
        function cost = coststopsigns(obj, xs)
        % Computes cost on max velocity through stop sign regions.
            max_stop_sign_v = 0;
            for i = 1:size(xs,1)
                if obj.scene.isCarFrontInStopRegion(xs(i,:))
                    max_stop_sign_v = max(max_stop_sign_v, abs(xs(i,4)));
                end
            end
            cost = obj.ws.stop_signs*max_stop_sign_v^2/2;
        end
        
    end
    
end