classdef TrajectoryExploreCEM < handle
% For the scenario given, executes the cross entropy method with parameters
% specifying knot states and the time length between them. 
% Vehicle dynamics restricted to driving straight.
% State:
%   X: position
%   v: velocity
% Control Input:
%   a: acceleration
%
% Knots are kappai = [dti, Xi, vi] where dti is the segment duration 
% between knot state [Xi-1, vi-1] and [Xi, vi]. 
% Parameters are a list of knots [kappa1, ..., kappaN]
% 
% Each knot has a knot_type associated with it which indicates which 
% parameters of the knot are sampled: 
%     accelerate: X and v  (dt is computed for constant acceleration from
%        previous knot)
%     cruise: X (v is v of previous knot; dt is computed for zero
%        acceleration from previous knot)
%     stop: none (X is a stopping position; v is 0; dt is computed for 
%        constant acceleration from previous knot
%     dwell: dt (X is X of previous knot; v is 0)
% 
% The acceleration is constant between knots.
        
    properties
        %%%% KNOBS 
        % disp for debug
        do_disp = true;
        %diagonal elements of Covariance matrix C for each knot type
        kC_INIT_KNOT_ACC = [3, 3]; 
        kC_INIT_KNOT_CRUISE = 3; 
        kC_INIT_KNOT_STOP = [];
        kC_INIT_KNOT_DWELL = 9;
        % Sampling lower and upper bounds.
        kDT_SAMPLING_BOUNDS = [.5, 10];
        kX_SAMPLING_BOUNDS = [-inf, inf];
        kv_SAMPLING_BOUNDS = [.5, 15];
        % The time discritization period.
        dt = .025;
        % Bounds to the 
        xBounds = [-inf, inf;
                   0   , 15];
        uBounds = [-4,4];
        xBounds_inc = [0, 0;
                       0, 0];
        uBounds_inc = [-0, 0]
        % Weights and other parameters describing cost function.
        ws_XYgoal = 0;
        ws_velocity = 5;
        ws_velocity_desired = 5;
        ws_tf = 0;
        ws_acceleration = 5;
        ws_velocityf = 0;
        % Parameters for the cross entropy method.
        Ncem_iters = 10
        opts_N = 200
        opts_rho = 0.1
        % for plotting the CEM trajectories
        do_plot = true;
        Ntraj_skip_for_plot = 10;
        %%%% END KNOBS 
        
        x0
        t0
        scene
        ax
        frames = [];
        
        NI = 1 % Number of Inputs (a)
        NS = 2 % Number of States (X, v)

        Nknots
        Nparams_per_knot = 3
        Nparams
        
        %knot type IDs
        kKNOT_ACC = 1
        kKNOT_CRUISE = 2
        kKNOT_STOP = 3
        kKNOT_DWELL = 4
        
        NLTL_automatas
        LTL_automatas % LTL autmata object
        costs % cost object
        constraints % constraint object
        
        Xstops % position of stop lines
        knot_types
        
        % variables facilitating the cross entropy method on segments of
        % the trajectory.
        Nknots_generated;
        Nknots_seen;
        knot_x0s;
        knot_t0s;
        knot_symbolic_states;
        cur_knot_num_attempts;
        max_knot_num_attempts

        cem_iter;
        
        % not all parameters are sampled
        sampled_params_index
        
        plot_cnt = 1;
    end
    
    
    methods
        function obj = TrajectoryExploreCEM(num_knots, scenario, LTL_automatas, proposition_handlers, Xstops)
            obj.Nknots = num_knots;
            obj.Nparams = obj.Nknots * obj.Nparams_per_knot;
            
            obj.x0 = scenario.x0;
            obj.t0 = 0;
            obj.scene = scenario;
            
            obj.Xstops = Xstops;
            
            if obj.do_plot
                figure
                obj.ax = gca;
            end
            
            assert(numel(LTL_automatas)==numel(proposition_handlers));
            obj.NLTL_automatas = numel(LTL_automatas);
            obj.LTL_automatas = LTL_automatas;
            


            % set CostFunction weights and parameters.
            if obj.ws_XYgoal ~= 0
                ws.XYgoal = obj.ws_XYgoal;
            end
            if obj.ws_velocity ~= 0
                ws.velocity = obj.ws_velocity;
                ws.velocity_desired = obj.ws_velocity_desired;
            end
            if obj.ws_tf ~= 0
                ws.tf = obj.ws_tf;
            end
            if obj.ws_acceleration ~= 0
                ws.acceleration = obj.ws_acceleration;
            end
            if obj.ws_velocityf ~= 0
                ws.velocityf = obj.ws_velocityf;
            end
            obj.costs = CostFunctions(scenario, ws);
            
            % ConstraintFunctions.m assumes full state and input --- i.e. 
            %   x = [X, Y, theta, v]', u = [a, psi]'
            cs.xBounds = [-inf*ones(4,1), inf*ones(4,1)];
            cs.uBounds = [-inf*ones(2,1), inf*ones(2,1)];
            for i = obj.Ncem_iters:-1:1
                cs.xBounds(1,:) = obj.xBounds(1,:) + (obj.Ncem_iters-i)*obj.xBounds_inc(1,:);
                cs.xBounds(2,:) = obj.xBounds(2,:) + (obj.Ncem_iters-i)*obj.xBounds_inc(2,:);
                
                cs.uBounds(1,:) = obj.uBounds(1,:) + (obj.Ncem_iters-i)*obj.uBounds_inc(1,:);
                
                if nargin > 2
                    obj.constraints(i).object = ConstraintFunctions(scenario, cs, obj.LTL_automatas, proposition_handlers);
                else
                    obj.constraints(i).object = ConstraintFunctions(scenario, cs);
                end
            end
            
            obj.Nknots_generated = 0;
            obj.Nknots_seen = 0;
            obj.knot_x0s = zeros(obj.Nknots + 1, length(obj.x0));
            obj.knot_x0s(1,:) = obj.x0;
            obj.knot_t0s = zeros(1, obj.Nknots + 1);
            obj.knot_t0s(1) = obj.t0;
            
            obj.knot_symbolic_states = zeros(obj.NLTL_automatas,obj.Nknots + 1);
            for i = 1:obj.NLTL_automatas
                obj.knot_symbolic_states(i,1) = obj.LTL_automatas{i}.start_state_index;
            end
            obj.cur_knot_num_attempts = 0;
            obj.max_knot_num_attempts = 20;
        end
        
        function best_params = runCE(obj, params, knot_types)%sampled_params_index)
        % params are of the form: [dt1, X1, v1; dt2, X2, v2; ....]
        %   where dti is the segment duration between knot locations 
        %   Xi is the position of the knot
        %   vi is the velocity of the knot
        % knot_types gives the types of each knot (accelerate, cruise,
        %   stop, or dwell)
            obj.knot_types = knot_types;
            opts.N = obj.opts_N;
            opts.rho = obj.opts_rho;
            
            % Set up sampling based on the knot_types.
            obj.sampled_params_index = [];
            C_diag = [];
            opts.lb = [];
            opts.ub = [];
            for i = 1:length(knot_types)
                switch obj.knot_types(i)
                    case obj.kKNOT_ACC
                        cur_sampled_params_index = [2,3];
                        cur_C_diag = obj.kC_INIT_KNOT_ACC;
                        cur_lb = [obj.kX_SAMPLING_BOUNDS(1), obj.kv_SAMPLING_BOUNDS(1)];
                        cur_ub = [obj.kX_SAMPLING_BOUNDS(2), obj.kv_SAMPLING_BOUNDS(2)];
                    case obj.kKNOT_CRUISE
                        cur_sampled_params_index = 2;
                        cur_C_diag = obj.kC_INIT_KNOT_CRUISE;
                        cur_lb = obj.kX_SAMPLING_BOUNDS(1);
                        cur_ub = obj.kX_SAMPLING_BOUNDS(2);
                    case obj.kKNOT_STOP
                        cur_sampled_params_index = [];
                        cur_C_diag = obj.kC_INIT_KNOT_STOP;
                        cur_lb = [];
                        cur_ub = [];
                    otherwise %obj.kKNOT_DWELL
                        cur_sampled_params_index = 1;
                        cur_C_diag = obj.kC_INIT_KNOT_DWELL;
                        cur_lb = obj.kDT_SAMPLING_BOUNDS(1);
                        cur_ub = obj.kDT_SAMPLING_BOUNDS(2);
                end
                            
                obj.sampled_params_index = [obj.sampled_params_index, cur_sampled_params_index + (i-1)*obj.Nparams_per_knot];
                C_diag = [C_diag, cur_C_diag];
                opts.lb = [opts.lb, cur_lb];
                opts.ub = [opts.ub, cur_ub];
            end
            opts.C = diag(C_diag);
            opts.iter = 1;
            opts.v = 1;
            opts.cFun = @obj.constraintFunc;
            obj.Nknots_generated = 0;           

            mu = params(obj.sampled_params_index)';
            for i = 1:obj.Ncem_iters
                obj.cem_iter = i;
                if obj.do_plot
                    % Clear axes for drawing new iteration's trajectories.
                    axes(obj.ax);
                    cla
                    hold on
                    obj.scene.drawFrame(obj.ax, [0,0], obj.x0(1:2), obj.x0(3))
                end
                [best_params, c, mu, C] = generativeCEM(@obj.costFunc, mu, opts);
                opts.C = C;
                
                disp([i, max(max(C))])
            end
            knots = obj.computeKnots(best_params);
            params = reshape(knots',1,numel(knots));
            best_params = params;
        end
        
        function cost = costFunc(obj, sampled_params)
        % Function executed during the cross entropy method for
        % computing the cost of the passed sampled_params.
            knots = obj.computeKnots(sampled_params);
            params = reshape(knots',1,numel(knots));
            
            [ts, xs, us] = obj.traj(params);
            
            % If costFunc called then a valid (constraint satisfying)
            % trajectory is found and will be kept. Plot if desired.
            if obj.do_plot
                obj.plot_cnt = obj.plot_cnt + 1;
                if mod(obj.plot_cnt,obj.Ntraj_skip_for_plot) == 0
                    % draw.
                    axes(obj.ax)
                    %plot(ts,xs(:,1),'blue',ts,xs(:,4),'red')
                    plot(xs(:,1),xs(:,4));
                    [~, knots_X, knots_v] = obj.paramsToKnots(params);
                    plot(knots_X, knots_v, 'o')
                    obj.frames = [obj.frames,getframe];
                    pause(.001)
                end
            end
            
            cost = obj.costs.compute(ts, xs, us);
        end
        function [is_satisfy, resample_from]= constraintFunc(obj, sampled_params)
        % Function executed during the cross entropy method for
        % computing if the constraints are satisfied for the passed 
        % sampled_params.
            if isempty(obj.constraints(obj.cem_iter))
                is_satisfy = true;
                resample_from = length(sampled_params);
            else
                knots = obj.computeKnots(sampled_params);
                params = reshape(knots',1,numel(knots));

                if obj.do_plot
                    disp(knots)
                end
                
                are_constraints_satisfied = true;
                %check monotonicity of X
                if  ~obj.isMonotonic([obj.x0(1), knots(:,2)']) || ~ obj.isGTMinDeltaTime(knots(:,1)')
                    are_constraints_satisfied = false;
                end
                while are_constraints_satisfied
%                   %%%%% NOTE--CARE! if full state can be saved, use this (e.g. proposition_handlers does not need to know time since an event occured):
%                     cur_param_range = obj.Nknots_generated*obj.Nparams_per_knot + 1:(obj.Nknots_seen + 1)*obj.Nparams_per_knot;
%                     cur_params = params(cur_param_range);
% 
%                     % computes trajectory from the previous knot to the
%                     % current one.
%                     [ts, xs, us] = obj.traj(cur_params, obj.knot_x0s(obj.Nknots_generated + 1,:), obj.knot_t0s(obj.Nknots_generated + 1));
%                     are_constraints_satisfied = obj.constraints.compute(ts, xs, us, obj.knot_symbolic_states(obj.Nknots_generated + 1));

                    %%%%% NOTE! otherwise use this (SAFER but less efficient)
                    cur_param_range = 1:(1 + obj.Nknots_seen)*obj.Nparams_per_knot;
                    cur_params = params(cur_param_range);

                    % build trajectory for current segment.
                    [ts, xs, us] = obj.traj(cur_params, obj.knot_x0s(1,:), obj.knot_t0s(1));
                    % check constraints.
                    are_constraints_satisfied = obj.constraints(obj.cem_iter).object.compute(ts, xs, us, obj.knot_symbolic_states(1));
                    if obj.do_disp
                        if are_constraints_satisfied
                            disp('SUCCEEDED!!')
                        else
                            disp('FAILED!!')
                        end
                    end
                    %%%%% END NOTE!
                    
                    if are_constraints_satisfied
                        obj.Nknots_generated = obj.Nknots_seen + 1;
                        obj.Nknots_seen = obj.Nknots_seen +1;
                                                
                        if obj.Nknots_generated == obj.Nknots
                            % full (all segments) trajectory is valid!
                            is_satisfy = true;
                            resample_from = length(sampled_params);
           
                            % set up for next sequence of knots.
                            obj.Nknots_generated = 0;
                            obj.Nknots_seen = 0;
                            obj.cur_knot_num_attempts = 0;
                            return
                        end
                        obj.knot_x0s(obj.Nknots_generated + 1,:) = xs(end,:);
                        obj.knot_t0s(obj.Nknots_generated + 1) = ts(end);
                        for j = 1:obj.NLTL_automatas
                            obj.knot_symbolic_states(j,obj.Nknots_generated + 1) = obj.LTL_automatas{j}.cur_state_index;
                        end
                    end
                
                end
                obj.cur_knot_num_attempts = obj.cur_knot_num_attempts + 1;

                if obj.do_disp
                    disp([obj.Nknots_generated, obj.Nknots_seen])
                end
                
                % grow segment backwards if too many failures over current
                % segment.
                if obj.cur_knot_num_attempts >= obj.max_knot_num_attempts 
                    if obj.Nknots_generated > 0
                        obj.Nknots_generated = obj.Nknots_generated - 1;
                    end
                    obj.cur_knot_num_attempts = 0;
                end
                
                is_satisfy = false;
                resample_from = length(intersect(obj.sampled_params_index, 1:obj.Nknots_generated*obj.Nparams_per_knot))+1;
            end
        end
        function [ts, xs, us] = traj(obj, params, x0, t0)
            % compute trajectory starting from initial state x0 and time t0
            % for passed params.
            if nargin <= 2
                x0 = obj.x0;
                t0 = 0;
            end
            
            [t_segs, knots_X, knots_v] = obj.paramsToKnots(params);
            tf = t0 + sum(t_segs);
            ts = t0:obj.dt:tf;
            
            
            % CostFunction.m and ConstraintFunction.m use full state i.e. 
            % x = [X, Y, theta, v];
            xs = zeros(length(ts),obj.NS);
            xs(:,2) = x0(2); % Y constant.
            xs(:,3) = x0(3); % theta constant.
            us = zeros(length(ts)-1,obj.NI);
            t_segs_cum = t0 + cumsum(t_segs);
            for i = 1:length(ts);
                [x,u] = obj.trajOft(t_segs_cum, knots_X, knots_v, ts(i), x0, t0);
                xs(i,[1,4]) = x;
                us(i) = u;
            end
        end
        
        function knots = computeKnots(obj, sampled_params)
            % computes the knots for the passed sampled_params. These
            % sampled_params depend on the sequence of knot_types.
            knots = zeros(obj.Nknots, obj.Nparams_per_knot);
            cur_param_index = 1;
            stop_counter = 1;
            for i = 1:obj.Nknots
                if i == 1
                    X0 = obj.x0(1);
                    v0 = obj.x0(4);
                else
                    X0 = knots(i-1,2);
                    v0 = knots(i-1,3);
                end
                switch obj.knot_types(i)
                    case obj.kKNOT_ACC
                        X1 = sampled_params(cur_param_index);
                        v1 = sampled_params(cur_param_index+1);
                        cur_param_index = cur_param_index + 2;
                        knots(i,1) = 2*(X1 - X0)/(v1 + v0);
                        knots(i,2) = X1;
                        knots(i,3) = v1;
                    case obj.kKNOT_CRUISE
                        X1 = sampled_params(cur_param_index);
                        v1 = v0;
                        cur_param_index = cur_param_index + 1;
                        knots(i,1) = (X1 - X0)/v1;
                        knots(i,2) = X1;
                        knots(i,3) = v1;
                    case obj.kKNOT_STOP
                        v1 = 0;
                        X1 = obj.Xstops(stop_counter);
                        knots(i,1) = 2*(X1 - X0)/(v1 + v0);
                        knots(i,2) = X1;
                        knots(i,3) = v1;
                        stop_counter = stop_counter + 1;
                    otherwise %obj.kKNOT_DWELL
                        DT1 = sampled_params(cur_param_index);
                        cur_param_index = cur_param_index + 1;
                        knots(i,1) = DT1;
                        assert(i>1);
                        knots(i,2) = knots(i-1,2);
                        knots(i,3) = 0;
                end
            end 
        end
        
        function [x, u] = trajOft(~, t_segs_cum, knots_X, knots_v, t, x0, t0)
            % computes the point on the trajectory at time t.
            i = find(t_segs_cum<t, 1, 'last')+1;
            if isempty(i)
                if t >= t0
                    i = 1;
                else
                    error('In function trajOft(), t=%f < t0', t);
                end
            end
            
            if i == 1
                % t0 = t0; % already t0
                X0 = x0(1);
                v0 = x0(4);
            else
                t0 = t_segs_cum(i-1);
                X0 = knots_X(i-1);
                v0 = knots_v(i-1);
            end
            t1 = t_segs_cum(i);
            X1 = knots_X(i); 
            v1 = knots_v(i);
            
            % u(t) = acceleration(t)     = c
            % v(t) = xs(2) = velocity(t) = c*(t-t0) + v0
            % X(t) = xs(1) = position(t) = c*(t-t0)^2/2 + v0*(t-t0) + X0
            % c such that X(t0) = X0, v(t0) = v0, X(t1) = X1 and v(t1) = v1
            %    assuming t1 = t0 + 2(X1-X0)/(v1+v0)
            if v1==v0
                c = 0;
            else
                c = (v1^2-v0^2)/(2*(X1 - X0));
            end
            u = c;
            x(2) = c*(t-t0) + v0;
            x(1) = c*(t-t0)^2/2 + v0*(t-t0) + X0;
        end
 
        function [t_segs, knots_X, knots_v] = paramsToKnots(obj,params)
        % Transforms params into knots.
            params = reshape(params, obj.Nparams_per_knot, numel(params)/obj.Nparams_per_knot)';
            t_segs = params(:,1)';
            knots_X = params(:,2)';
            knots_v = params(:,3)';
            
%             t_segs(t_segs < .1) = .1; % hack to make all negative time segments be zero time.
%             knots_v(knots_v < .1) = .1; % hack to make all negative velocities be zero velocity.
        end
        
        function is_monotonic = isMonotonic(~,list)
            % checks if list is monotonically increasing---i.e. never decreasing.
            is_monotonic = all(diff(list)>=0);
        end
        
        function is_gt_min_dt = isGTMinDeltaTime(obj,dts)
            % checks if dts are sufficiently long
            is_gt_min_dt = all(dts > 3*obj.dt);
        end
        
    end
end