when-to-switch

1
// cppimport
2
#include <pybind11/pybind11.h>
3
#include <pybind11/stl.h>
4
#include <pybind11/stl_bind.h>
5
#include <vector>
6
#include <queue>
7
#include <cmath>
8
#include <set>
9
#include <map>
10
#include <list>
11
#define INF 1000000000
12
namespace py = pybind11;
13

14
struct Node {
15
    Node(int _i = INF, int _j = INF, int _g = 0, int _h = 0) : i(_i), j(_j), g(_g), h(_h), f(_g+_h){}
16
    int i;
17
    int j;
18
    int g;
19
    int h;
20
    int f;
21
    bool operator<(const Node& other) const
22
    {
23
        return this->f < other.f or
24
               (this->f == other.f and (this->g < other.g or
25
                                       (this->g == other.g and (this->i < other.i or
26
                                                               (this->i == other.i and this->j < other.j)))));
27
    }
28
    bool operator>(const Node& other) const
29
    {
30
        return this->f > other.f or
31
               (this->f == other.f and (this->g > other.g or
32
                                       (this->g == other.g and (this->i > other.i or
33
                                                               (this->i == other.i and this->j > other.j)))));
34
    }
35
    bool operator==(const Node& other) const
36
    {
37
        return this->i == other.i and this->j == other.j;
38
    }
39
    bool operator==(const std::pair<int, int> &other) const
40
    {
41
        return this->i == other.first and this->j == other.second;
42
    }
43
};
44

45
class planner {
46
    std::set<std::pair<int,int>> obstacles;
47
    std::set<std::pair<int,int>> other_agents;
48
    std::set<std::pair<int,int>> bad_actions;
49
    std::priority_queue<Node, std::vector<Node>, std::greater<Node>> OPEN;
50
    std::map<std::pair<int, int>, std::pair<int, int>> CLOSED;
51
    std::pair<int, int> start;
52
    std::pair<int, int> desired_position;
53
    std::pair<int, int> goal;
54
    Node best_node;
55
    int max_steps;
56
    inline int h(std::pair<int, int> n)
57
    {
58
        return std::abs(n.first - goal.first) + std::abs(n.second - goal.second);
59
    }
60
    std::vector<std::pair<int,int>> get_neighbors(std::pair<int, int> node)
61
    {
62
        std::vector<std::pair<int,int>> neighbors;
63
        std::vector<std::pair<int,int>> deltas = {{0,1},{1,0},{-1,0},{0,-1}};
64
        for(auto d:deltas)
65
        {
66
            std::pair<int,int> n(node.first + d.first, node.second + d.second);
67
            if(obstacles.count(n) == 0)
68
                neighbors.push_back(n);
69
        }
70
        return neighbors;
71
    }
72
    void compute_shortest_path()
73
    {
74
        Node current;
75
        int steps = 0;
76
        while(!OPEN.empty() and steps < max_steps and !(current == goal))
77
        {
78
            current = OPEN.top();
79
            OPEN.pop();
80
            if(current.h < best_node.h)
81
                best_node = current;
82
            steps++;
83
            for(auto n: get_neighbors({current.i, current.j})) {
84
                if (CLOSED.find(n) == CLOSED.end() and other_agents.find(n) == other_agents.end()) {
85
                    OPEN.push(Node(n.first, n.second, current.g + 1, h(n)));
86
                    CLOSED[n] = {current.i, current.j};
87
                }
88
            }
89
        }
90
    }
91
    void reset()
92
    {
93
        CLOSED.clear();
94
        OPEN = std::priority_queue<Node, std::vector<Node>, std::greater<Node>>();
95
        Node s = Node(start.first, start.second, 0, h(start));
96
        OPEN.push(s);
97
        best_node = s;
98
    }
99
public:
100
    planner(int steps=10000) {max_steps = steps;}
101
    void update_obstacles(const std::list<std::pair<int, int>>& _obstacles,
102
                          const std::list<std::pair<int, int>>& _other_agents,
103
                          std::pair<int, int> cur_pos)
104
    {
105
        for(auto o:_obstacles)
106
            obstacles.insert({cur_pos.first + o.first, cur_pos.second + o.second});
107
        other_agents.clear();
108
        for(auto o:_other_agents)
109
            other_agents.insert({cur_pos.first + o.first, cur_pos.second + o.second});
110
    }
111
    void update_path(std::pair<int, int> s, std::pair<int, int> g)
112
    {
113
        if(desired_position.first < INF and (desired_position.first != s.first or desired_position.second != s.second)) {
114
            bad_actions.insert(desired_position);
115
            if (start.first == s.first and start.second == s.second)
116
                for (auto bad_a: bad_actions)
117
                    other_agents.insert(bad_a);
118
        }
119
        else
120
            bad_actions.clear();
121
        start = s;
122
        goal = g;
123
        reset();
124
        compute_shortest_path();
125
    }
126
    std::list<std::pair<int, int>> get_path(bool use_best_node = true)
127
    {
128
        std::list<std::pair<int, int>> path;
129
        std::pair<int, int> next_node(INF,INF);
130
        if(CLOSED.find(goal) != CLOSED.end())
131
            next_node = goal;
132
        else if(use_best_node)
133
            next_node = {best_node.i, best_node.j};
134
        if(next_node.first < INF and (next_node.first != start.first or next_node.second != start.second))
135
        {
136
            while (CLOSED[next_node] != start) {
137
                path.push_back(next_node);
138
                next_node = CLOSED[next_node];
139
            }
140
            path.push_back(next_node);
141
            path.push_back(start);
142
            path.reverse();
143
        }
144
        desired_position = next_node;
145
        return path;
146
    }
147
    std::pair<std::pair<int, int>, std::pair<int, int>> get_next_node(bool use_best_node = true)
148
    {
149
        std::pair<int, int> next_node(INF, INF);
150
        if(CLOSED.find(goal) != CLOSED.end())
151
            next_node = goal;
152
        else if(use_best_node)
153
            next_node = {best_node.i, best_node.j};
154
        if(next_node.first < INF and (next_node.first != start.first or next_node.second != start.second))
155
            while (CLOSED[next_node] != start)
156
                next_node = CLOSED[next_node];
157
        if(next_node == start)
158
            next_node = {INF, INF};
159
        desired_position = next_node;
160
        return {start, next_node};
161
    }
162
};
163

164
PYBIND11_MODULE(planner, m) {
165
    py::class_<planner>(m, "planner")
166
            .def(py::init<int>())
167
            .def("update_obstacles", &planner::update_obstacles)
168
            .def("update_path", &planner::update_path)
169
            .def("get_path", &planner::get_path)
170
            .def("get_next_node", &planner::get_next_node);
171
}
172

173
/*
174
<%
175
setup_pybind11(cfg)
176
%>
177
*/
178