classdef BuildLattice < handle
    % Builds a grid of X (position or arc length) and v (velocity) points
    % and connects them assuming constant acceleration. 
    properties
        x0
        dgrid
        grid_end
        acc_max
        dt_max
        position_adjacency_factor
        
        state_knots
        state_knots_index
        nonstop_position_ranges % e.g. at crosswalk.
        stop_position_ranges % e.g. at stopsign.
        stop_max_v = 1.1;
        
        X_grid
        v_grid
        grid_size
        
        is_connections
        is_vertices
    end
    methods
        % x0 is the initial state.
        %   x0(1) = X0 (initial position)
        %   x0(4) = v0 (initial velocity)
        %   x0(2) and x0(3) are Y position and angle that are not used here
        %   for driving straight.
        % dgrid is the grid spacing.
        %   dgrid(1) = dX
        %   dgrid(2) = dVelocity
        % grid_end is upper bounds on the grid.
        %   grid_end(1) = Xend
        %   grid_end(2) = vend
        % acc_max is max acceleration an edge can have.
        % dt_max is the max time an edge can have.
        % position_adjacency_factor (horrible name). It is the maximum 
        %   number of dX between vertices. 
        %   e.g. position_adjacency_factor = 3 implies (X1, v1) could
        %   possibly connect to (X4, v1) but not (X5, v1).
        % state_knots are list of position, velocity pairs that a
        %   trajectory must go through. 
        %   e.g. state_knot(i,:) = [10, 0] means the vertex at 
        %       X = 10, v = 0 must exist but all other vertices at X = 10
        %       do not.
        % nonstop_position_ranges (again horrible name) are the positions
        %   that the vehicle is not allowed to stop at.
        %   e.g. nonstop_position_ranges(i,:) = [5, 9] means that any
        %   vertex (X_i, 0) where 5 <= X_i <= 9 is not allowed 
        function obj = BuildLattice(x0, dgrid, grid_end, acc_max, dt_max, position_adjacency_factor, state_knots, stop_position_ranges, nonstop_position_ranges)
            obj.x0 = x0;
            obj.dgrid = dgrid;
            obj.grid_end = grid_end;
            obj.acc_max = acc_max;
            obj.dt_max = dt_max;
            obj.position_adjacency_factor = 1;
            if nargin >= 6
                obj.position_adjacency_factor = position_adjacency_factor;
            end
            
            obj.makeRawGrid();
            obj.makeRawConnections();
            obj.setStateKnots([]);
            if nargin >= 7
                obj.setStateKnots(state_knots);
            end
            obj.stop_position_ranges = [];
            if nargin >= 8
                obj.stop_position_ranges = stop_position_ranges;
            end
            obj.nonstop_position_ranges = [];
            if nargin >= 9
                obj.nonstop_position_ranges = nonstop_position_ranges;
            end
            obj.pruneVerticesAndConnections();
        end

        % state knots are position, velocity pairs that the trajectory must
        % traverse through. Most are 0 velocity positions for stopping.
        %
        % state knots are in the form
        % state_knots = [knot1_X, knot1_v;
        %                knot2_X, knot2_v;
        %                  ...  ,   ...  ;
        %                knotn_X, knotn_v]
        % The initial state is treated as a state knot and should not be
        % passed as a state_knot.
        function setStateKnots(obj, state_knots)
            obj.state_knots = [obj.x0(1), obj.x0(4);   
                               state_knots];
            obj.state_knots_index = zeros(size(obj.state_knots,1),2);
            for i = 1:size(obj.state_knots,1)
                [index_X, index_v] = obj.nearestVertexToState(obj.state_knots(i,:));
                obj.state_knots_index(i,:) = [index_X, index_v];
            end
        end
        
        % grids up the space from the initial state position to the given
        % end state and 0 velocity to the max velocity (both specified by
        % grid_end). The granularity is given by dgrid. All vertices are
        % set to be valid at this stage.
        function makeRawGrid(obj)
            obj.X_grid = obj.x0(1):obj.dgrid(1):obj.grid_end(1);
            obj.v_grid = 0:obj.dgrid(2):obj.grid_end(2);
            obj.grid_size = [length(obj.v_grid) length(obj.X_grid)];
            obj.is_vertices = true(obj.grid_size);
        end
        
        % Builds is_connections matrix.
        % connections are always of the form (X_i,v_j) -> (X_{i+l},v_k)  where
        % connections dimensions are in order: i,j,k,l
        % connections is num_X x num_v x num_v.
        % All vertices are assumed to be valid at this stage.
        function makeRawConnections(obj)
            obj.is_connections = false([obj.grid_size(2), obj.grid_size(1), obj.grid_size(1), obj.position_adjacency_factor]);
            for j = 1:obj.grid_size(1)
                % Loops through from velocities.
                for k = 1:obj.grid_size(1)
                    % Loops through to velocities.
                    for l = 1:obj.position_adjacency_factor
                        % Loops through position adjacency distance.
                        state1 = [obj.X_grid(1), obj.v_grid(j)];
                        state2 = [obj.X_grid(1+l), obj.v_grid(k)];

                        [dt, a] = obj.dtAndAccForStateTransfer(state1, state2);
                        if abs(a) < obj.acc_max && dt < obj.dt_max 
                            obj.is_connections(1:end-l,j,k,l) = true;
                        end
                    end
                end
            end
        end
        
        % Prunes the vertices and connections that are not possible due to
        % the restriction that trajectories must flow through state_knots.
        % First pass is forward in time from each state_knot up to the next
        % state knot. Second pass is backward in time from each state_knot
        % to the prior state_knot.
        % TODO(remove connections that jump a state_knot).
        function pruneVerticesAndConnections(obj)
            for i = 1:size(obj.state_knots,1)
                %loop forward through state_knots.
                
                stop_index_X = obj.grid_size(2);
                if i < size(obj.state_knots,1)
                    stop_index_X = obj.state_knots_index(i+1,1);
                end
                for j = obj.state_knots_index(i,1):stop_index_X
                    % loop through positions from current state knot to next
                    % state knot or the end.
                    
                    
                    for k = 1:size(obj.nonstop_position_ranges,1)
                        if obj.nonstop_position_ranges(k,1) <= obj.X_grid(j) ...
                                && obj.X_grid(j) <= obj.nonstop_position_ranges(k,2)
                            disp('HERE!!')
                            % can't stop here!
                            obj.is_vertices(1,j) = false; 
                            break
                        end
                    end
                    
                    for k = 1:size(obj.stop_position_ranges,1)
                        if obj.stop_position_ranges(k,1) <= obj.X_grid(j) ...
                                && obj.X_grid(j) <= obj.stop_position_ranges(k,2)
                            
                            % can't go fast here!
%                             bad_range = find(obj.v_grid > obj.stop_max_v);
                            obj.is_vertices(obj.v_grid > obj.stop_max_v,j) = false; 
                            break
                        end
                    end
                    
                    if j == obj.state_knots_index(i,1)
                        obj.is_vertices(:,j) = false;
                        obj.is_vertices(obj.state_knots_index(i,2),j) = true;
                        
                        % remove connections that jump the state knot position
                        for k = 1:obj.grid_size(1)
                        % loop through velocities at this position
                            if ~obj.is_vertices(k,j)
                                for l = 1:obj.position_adjacency_factor
                                    if j-l >= 1
                                        obj.is_connections(j-l,:,k,l) = false;
                                    end
                                end
                            end
                        end
                    else
                        % invalidate vertices at this postion with no connections to it.
                        for k = 1:obj.grid_size(1)
                            % loop through velocities at this position
                            if obj.is_vertices(k,j)
                                obj.is_vertices(k,j) = false;
                                for l = 1:obj.position_adjacency_factor
                                    if j-l >= obj.state_knots_index(i,1)
                                        if sum(obj.is_connections(j-l,:,k,l))>0
                                            obj.is_vertices(k,j) = true;
                                            break;
                                        end
                                    end
                                end
                            end
                        end
                    end

                    % remove all connections that jump a state_knot.
                    if j > obj.state_knots_index(i,1) && j < obj.state_knots_index(i,1) + obj.position_adjacency_factor% + 1
                        for l = 2:obj.position_adjacency_factor
                            if j-l >= 1 && j-l < obj.state_knots_index(i,1)
                                obj.is_connections(j-l,:,:,l) = false;
                            end
                        end    
                    end
                        
                    % remove connections from invalid vertices.
                    for k = 1:obj.grid_size(1)
                        % loop through velocities at this position
                        if ~obj.is_vertices(k,j)
                            % remove all connections from an invalid vertex
                            obj.is_connections(j,k,:,:) = false;
                            % remove all connections to an invalid vertex
                            for l = 1:obj.position_adjacency_factor
                                if j-l >= 1
                                    obj.is_connections(j-l,:,k,l) = false;
                                end
                            end    
                        end
                    end
                end
            end
            
            index_loop = [obj.state_knots_index(:,1)',length(obj.X_grid)];
            disp(index_loop)
            for i = length(index_loop):-1:2
                %loop backward through state_knots.
               
                stop_index_X = index_loop(i-1,1);
                for j = index_loop(i):-1:stop_index_X
                    % loop through positions from current state knot to next
                    % state knot or the end.
                    
                    % invalidate vertices at this postion with no connections from it.
                    if j < index_loop(i)
                        for k = 1:obj.grid_size(1)
                            % loop through velocities at this position
                            if obj.is_vertices(k,j)
                                obj.is_vertices(k,j) = false;
                                for l = 1:obj.position_adjacency_factor
                                    if j+l <= obj.grid_size(2)
                                        if sum(obj.is_connections(j,k,:,l))>0
                                            obj.is_vertices(k,j) = true;
                                            break;
                                        end
                                    end
                                end
                            end
                        end
                    end
                    
                    % remove connections from invalid vertices.
                    for k = 1:obj.grid_size(1)
                        % loop through velocities at this position
                        if ~obj.is_vertices(k,j)
%                             % remove all connections from an invalid vertex
%                             obj.is_connections(j,k,:,:) = false;
                            % remove all connections to an invalid vertex
                            for l = 1:obj.position_adjacency_factor
                                if j-l >= 1
                                    obj.is_connections(j-l,:,k,l) = false;
                                end
                            end    
                        end
                    end
                end
            end
        end
        
        function [index_X, index_v] = nearestVertexToState(obj, x)
            [~, index_X] = min(abs(x(1) - obj.X_grid));
            [~, index_v] = min(abs(x(2) - obj.v_grid));
        end
        
        % Computes constant acceleration and time to transfer the system from state
        % x1 to x2 where states are (X, v)---i.e. just position and
        % velocity.
        function [dt, a] = dtAndAccForStateTransfer(~, x1, x2)
            assert(x1(1) ~= x2(1)) % can't be at same position
            
            a = (x2(2)^2 - x1(2)^2)/(2*(x2(1) - x1(1)));
            if x2(2) == 0 && x2(1) == 0
                dt = inf;
            else
                dt = 2*(x2(1) - x1(1))/(x2(2) + x1(2));
            end
        end
        
        function plotVertices(obj, ax, offset)
            if nargin >= 2
                axes(ax);
            end
            if nargin < 3
                offset = 0;
            end
            [pts_X, pts_v] = meshgrid(obj.X_grid,obj.v_grid);
            plot(pts_X,pts_v + offset,'bo');
            for i = 1:obj.grid_size(2)%-1
                % Loops through positions.
                for j = 1:obj.grid_size(1)
                    % Loops through from velocities.
                    if obj.is_vertices(j,i)
                        plot(obj.X_grid(i), obj.v_grid(j) + offset, 'ro')
                    end
                end
            end
        end
        function plotEdges(obj, ax, offset)
            if nargin >= 2
                axes(ax);
            end
            if nargin < 3
                offset = 0;
            end
            for i = 1:obj.grid_size(2)-1
                % Loops through positions.
                for j = 1:obj.grid_size(1)
                    % Loops through from velocities.
                    for k = 1:obj.grid_size(1)
                        % Loops through to velocities.
                        for l = 1:obj.position_adjacency_factor
                            % Loops through position adjacency distance.
                            if obj.is_connections(i,j,k,l) && i+l <= obj.grid_size(2)
                                list1 = [obj.X_grid(i), obj.X_grid(i+l)];
                                list2 = [obj.v_grid(j) + offset, obj.v_grid(k) + offset];
                                plot(list1,list2,'b')
                            end
                        end
                    end
                end
            end
        end
        function plotGrid(obj, ax)
            if nargin >= 2
                axes(ax);
            end
            obj.plotVertices();
            obj.plotEdges();
        end
    end
    
    
end