1use std::cmp::min;
2
3use indexmap::IndexSet;
4
5use crate::coord::Coord;
6pub use crate::coord::Direction;
7
8use rand::rngs::StdRng;
9use rand::{Rng, SeedableRng};
10
11#[cfg(feature = "display")]
12use sfml::graphics::Color;
13#[cfg(feature = "display")]
14pub use {sfml::graphics::RenderWindow, sfml::window::Style};
15
16pub struct Snake {
18 pub(super) snake: Vec<Coord>,
19 empty: IndexSet<Coord>,
20 dir: Direction,
21 food: Coord,
22 pub(super) size: u8,
23 rng: StdRng,
24 #[cfg(feature = "display")]
25 pub(super) display: Option<RenderWindow>,
26}
27
28impl Snake {
29 pub fn new(seed: u64, size: u8) -> Snake {
43 let snake_first = Coord {
44 x: size / 2,
45 y: size / 2 - 1,
46 };
47 let snake_second = Coord {
48 x: size / 2 - 1,
49 y: size / 2 - 1,
50 };
51
52 let snake = vec![snake_first, snake_second];
53
54 let mut empty = IndexSet::new();
55
56 for x in 0..size {
57 for y in 0..size {
58 empty.insert(Coord { x, y });
59 }
60 }
61
62 empty.remove(&snake_first);
63 empty.remove(&snake_second);
64
65 let mut s = Snake {
66 snake,
67 empty,
68 dir: Direction::Right,
69 food: Coord { x: 0, y: 0 },
70 size,
71 rng: SeedableRng::seed_from_u64(seed),
72 #[cfg(feature = "display")]
73 display: None,
74 };
75
76 s.gen_food();
77 s
78 }
79
80 #[cfg(feature = "display")]
97 pub fn new_display(seed: u64, size: u8, display: Option<RenderWindow>) -> Snake {
98 let mut s = Snake::new(seed, size);
99 s.display = display;
100
101 s.init_display();
102 s.draw_square(s.food, Color::GREEN);
103 s.display();
104 s
105 }
106
107 pub fn length(&self) -> usize {
109 self.snake.len()
110 }
111
112 pub fn current_direction(&self) -> Direction {
114 self.dir
115 }
116
117 pub fn walls(&self) -> Vec<f32> {
120 let mut walls = Vec::with_capacity(8);
121
122 let size = self.size - 1;
123 let head = self.snake.first().unwrap();
124
125 walls.push(head.x as f32 / self.size as f32);
126 walls.push(min(head.x, head.y) as f32 / self.size as f32);
127 walls.push(head.y as f32 / self.size as f32);
128 walls.push(min(size - head.x, head.y) as f32 / self.size as f32);
129 walls.push((size - head.x) as f32 / self.size as f32);
130 walls.push(min(size - head.x, size - head.y) as f32 / self.size as f32);
131 walls.push((size - head.y) as f32 / self.size as f32);
132 walls.push(min(head.x, size - head.y) as f32 / self.size as f32);
133
134 walls
135 }
136
137 fn snake_in_dir(&self, dir: Direction, dir2: Direction) -> f32 {
138 let mut c = *self.snake.first().unwrap();
139
140 for i in 1..self.size {
141 c = match c + dir {
142 Err(_) => return 1.0,
143 Ok(val) => val,
144 };
145 c = match c + dir2 {
146 Err(_) => return 1.0,
147 Ok(val) => val,
148 };
149
150 if !c.in_bounds(self.size) {
151 return 1.0;
152 }
153 if !self.empty.contains(&c) {
154 return i as f32 / self.size as f32;
155 }
156 }
157
158 1.0
159 }
160
161 pub fn snake(&self) -> Vec<f32> {
164 let mut snake = Vec::with_capacity(8);
165
166 snake.push(self.snake_in_dir(Direction::Left, Direction::Center));
167 snake.push(self.snake_in_dir(Direction::Up, Direction::Left));
168 snake.push(self.snake_in_dir(Direction::Up, Direction::Center));
169 snake.push(self.snake_in_dir(Direction::Up, Direction::Right));
170 snake.push(self.snake_in_dir(Direction::Right, Direction::Center));
171 snake.push(self.snake_in_dir(Direction::Down, Direction::Right));
172 snake.push(self.snake_in_dir(Direction::Down, Direction::Center));
173 snake.push(self.snake_in_dir(Direction::Down, Direction::Left));
174
175 snake
176 }
177
178 pub fn food(&self) -> Vec<f32> {
181 let head = self.snake.first().unwrap();
182 let food = self.food;
183
184 let mut dir = vec![1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0];
185
186 if head.y == food.y {
187 if head.x < food.x {
188 dir[4] = (food.x - head.x) as f32 / self.size as f32;
189 } else {
190 dir[0] = (head.x - food.x) as f32 / self.size as f32;
191 }
192 } else if head.x == food.x {
193 if head.y < food.y {
194 dir[2] = (food.y - head.y) as f32 / self.size as f32;
195 } else {
196 dir[6] = (head.y - food.y) as f32 / self.size as f32;
197 }
198 } else if head.x + head.y == food.x + food.y {
199 if head.x < food.x {
200 dir[3] = (food.x - head.x) as f32 / self.size as f32;
201 } else {
202 dir[7] = (head.x - food.x) as f32 / self.size as f32;
203 }
204 } else if (head.x as i16 - head.y as i16).abs() == (food.x as i16 - food.y as i16).abs() {
205 if head.x < food.x {
206 dir[5] = (food.x - head.x) as f32 / self.size as f32;
207 } else {
208 dir[1] = (head.x - food.x) as f32 / self.size as f32;
209 }
210 }
211
212 dir
213 }
214
215 fn alive(&self) -> bool {
218 let mut snake = self.snake.iter();
219 let head = snake.next().unwrap();
220
221 if head.x >= self.size || head.y >= self.size {
222 return false;
223 }
224
225 self.empty.contains(head)
226 }
227
228 fn found_food(&self) -> bool {
230 let head = self.snake.first().unwrap();
231 *head == self.food
232 }
233
234 fn gen_food(&mut self) -> bool {
236 if self.empty.is_empty() {
237 return false;
238 }
239
240 self.food = *self
241 .empty
242 .get_index(self.rng.gen_range(0, self.empty.len()))
243 .unwrap();
244
245 #[cfg(feature = "display")]
246 self.draw_square(self.food, Color::GREEN);
247
248 true
249 }
250
251 pub fn turn(&mut self, dir: Direction) -> bool {
260 match self.dir {
261 Direction::Left => {
262 if dir == Direction::Up || dir == Direction::Down {
263 self.dir = dir;
264 }
265 }
266 Direction::Right => {
267 if dir == Direction::Up || dir == Direction::Down {
268 self.dir = dir;
269 }
270 }
271 Direction::Up => {
272 if dir == Direction::Left || dir == Direction::Right {
273 self.dir = dir;
274 }
275 }
276 Direction::Down => {
277 if dir == Direction::Left || dir == Direction::Right {
278 self.dir = dir;
279 }
280 }
281 Direction::Center => panic!("Direction can't be center"),
282 }
283
284 let curr_pos = *self.snake.first().unwrap();
285
286 let new_head = match self.dir {
287 Direction::Left => {
288 if curr_pos.x == 0 {
289 return false;
290 }
291
292 (curr_pos + self.dir).unwrap()
293 }
294 Direction::Right => {
295 if curr_pos.x == self.size - 1 {
296 return false;
297 }
298
299 (curr_pos + self.dir).unwrap()
300 }
301 Direction::Up => {
302 if curr_pos.y == 0 {
303 return false;
304 }
305
306 (curr_pos + self.dir).unwrap()
307 }
308 Direction::Down => {
309 if curr_pos.y == self.size - 1 {
310 return false;
311 }
312
313 (curr_pos + self.dir).unwrap()
314 }
315 Direction::Center => panic!("Direction can't be center"),
316 };
317
318 self.snake.insert(0, new_head);
319
320 if !self.alive() {
321 return false;
322 }
323
324 self.empty.remove(&new_head);
325
326 if !self.found_food() {
327 let tail = self.snake.pop().unwrap();
328 self.empty.insert(tail);
329 #[cfg(feature = "display")]
330 self.draw_square(tail, Color::BLACK);
331 } else if !self.gen_food() {
332 return false;
333 }
334
335 #[cfg(feature = "display")]
336 {
337 let head = *self.snake.first().unwrap();
338 self.draw_square(head, Color::WHITE);
339
340 self.display();
341 }
342
343 true
344 }
345}
346
347#[cfg(feature = "display")]
349impl Drop for Snake {
350 fn drop(&mut self) {
351 if let Some(d) = self.display.as_mut() {
352 d.close();
353 }
354 }
355}
356
357#[cfg(test)]
358mod tests {
359
360 use super::*;
361
362 #[test]
363 fn test_alive() {
364 let mut test = Snake::new(0, 10);
365 test.snake.insert(0, Coord { x: 6, y: 4 });
366 assert!(test.alive());
367 test.snake = vec![Coord { x: 4, y: 10 }];
368 assert!(!test.alive());
369 test.snake = vec![
370 Coord { x: 5, y: 4 },
371 Coord { x: 4, y: 4 },
372 Coord { x: 4, y: 3 },
373 Coord { x: 5, y: 3 },
374 Coord { x: 5, y: 4 },
375 ];
376 assert!(!test.alive());
377 }
378
379 #[test]
380 fn test_found_food() {
381 let mut test = Snake::new(0, 10);
382 test.food = Coord { x: 0, y: 2 };
383 assert!(!test.found_food());
384 test.snake = vec![Coord { x: 0, y: 2 }];
385 assert!(test.found_food());
386 }
387
388 #[test]
389 fn test_snake() {
390 let mut test = Snake::new(0, 10);
391 test.food = Coord { x: 0, y: 2 };
392 assert!(test.turn(Direction::Up));
393 assert!(test.turn(Direction::Center));
394
395 assert!(test.turn(Direction::Left));
396 for _ in 0..4 {
397 assert!(test.turn(Direction::Center));
398 }
399 assert_eq!(test.length(), 3);
400 test.food = Coord { x: 9, y: 5 };
401
402 assert!(test.turn(Direction::Down));
403 for _ in 0..2 {
404 assert!(test.turn(Direction::Center));
405 }
406
407 assert!(test.turn(Direction::Right));
408 for _ in 0..8 {
409 assert!(test.turn(Direction::Center));
410 }
411 assert_eq!(test.length(), 4);
412 }
413
414 #[test]
415 fn test_walls() {
416 let test = Snake::new(0, 10);
417 assert_eq!(test.walls(), vec![0.5, 0.4, 0.4, 0.4, 0.4, 0.4, 0.5, 0.5]);
418 }
419
420 #[test]
421 fn test_snake_dis() {
422 let test = Snake::new(0, 10);
423 assert_eq!(test.snake(), vec![0.1, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]);
424 }
425
426 #[test]
427 fn test_food() {
428 let mut test = Snake::new(0, 10);
429 assert_eq!(test.food(), vec![1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]);
430 test.food = Coord { x: 8, y: 4 };
431 assert_eq!(test.food(), vec![1.0, 1.0, 1.0, 1.0, 0.3, 1.0, 1.0, 1.0]);
432 }
433}