package day21

import (
	"fmt"

	"gitea.paas.celticinfo.fr/oabrivard/aoc2023/utils"
)

type Coord struct {
	row int
	col int
}

func BFS1(matrix [][]byte, root Coord, maxSteps int) int {
	height := len(matrix)
	width := len(matrix[0])
	queue := utils.NewQueue[*Coord]()

	matrix[root.row][root.col] = 'O'
	queue.Enqueue(&root)

	moves := 0
	for moves < maxSteps {
		//fmt.Println(moves + 1)
		newQueue := utils.NewQueue[*Coord]()

		for queue.HasElement() {
			coord := queue.Dequeue()
			matrix[coord.row][coord.col] = '.'

			if coord.row > 0 && matrix[coord.row-1][coord.col] != '#' {
				newCoord := Coord{coord.row - 1, coord.col}
				newQueue.Enqueue(&newCoord)
				matrix[newCoord.row][newCoord.col] = 'O'
			}

			if coord.row < height-1 && matrix[coord.row+1][coord.col] != '#' {
				newCoord := Coord{coord.row + 1, coord.col}
				newQueue.Enqueue(&newCoord)
				matrix[newCoord.row][newCoord.col] = 'O'
			}

			if coord.col > 0 && matrix[coord.row][coord.col-1] != '#' {
				newCoord := Coord{coord.row, coord.col - 1}
				newQueue.Enqueue(&newCoord)
				matrix[newCoord.row][newCoord.col] = 'O'
			}

			if coord.col < width-1 && matrix[coord.row][coord.col+1] != '#' {
				newCoord := Coord{coord.row, coord.col + 1}
				newQueue.Enqueue(&newCoord)
				matrix[newCoord.row][newCoord.col] = 'O'
			}
		}

		queue = newQueue
		moves++
	}

	// count 'O' in matrix
	count := 0
	for r := range matrix {
		for _, b := range matrix[r] {
			if b == 'O' {
				count++
			}
		}
	}

	return count
}

func redrawFromCache(matrix [][]byte, cacheQueue utils.Queue[*Coord]) {
	for _, coord := range cacheQueue.GetData() {
		if coord != nil {
			matrix[coord.row][coord.col] = 'O'
		}
	}
}

func BFS2(matrix [][]byte, root Coord, maxSteps int) int {
	height := len(matrix)
	width := len(matrix[0])
	queue := utils.NewQueue[*Coord]()

	cache := make([][]utils.Queue[*Coord], height)
	for r, row := range matrix {
		cache[r] = make([]utils.Queue[*Coord], width)
		for c := range row {
			cache[r][c] = utils.NewQueue[*Coord]()
		}
	}

	matrix[root.row][root.col] = 'O'
	queue.Enqueue(&root)

	moves := 0
	for moves < maxSteps {
		fmt.Println(moves + 1)
		newQueue := utils.NewQueue[*Coord]()

		for queue.HasElement() {
			coord := queue.Dequeue()

			if cache[coord.row][coord.col].HasElement() {

				matrix[coord.row][coord.col] = '.'
				redrawFromCache(matrix, cache[coord.row][coord.col])

			} else {
				cacheQueue := utils.NewQueue[*Coord]()
				matrix[coord.row][coord.col] = '.'

				if coord.row > 0 && matrix[coord.row-1][coord.col] != '#' {
					newCoord := Coord{coord.row - 1, coord.col}
					newQueue.Enqueue(&newCoord)
					matrix[newCoord.row][newCoord.col] = 'O'
					cacheQueue.Enqueue(&newCoord)
				}

				if coord.row < height-1 && matrix[coord.row+1][coord.col] != '#' {
					newCoord := Coord{coord.row + 1, coord.col}
					newQueue.Enqueue(&newCoord)
					matrix[newCoord.row][newCoord.col] = 'O'
					cacheQueue.Enqueue(&newCoord)
				}

				if coord.col > 0 && matrix[coord.row][coord.col-1] != '#' {
					newCoord := Coord{coord.row, coord.col - 1}
					newQueue.Enqueue(&newCoord)
					matrix[newCoord.row][newCoord.col] = 'O'
					cacheQueue.Enqueue(&newCoord)
				}

				if coord.col < width-1 && matrix[coord.row][coord.col+1] != '#' {
					newCoord := Coord{coord.row, coord.col + 1}
					newQueue.Enqueue(&newCoord)
					matrix[newCoord.row][newCoord.col] = 'O'
					cacheQueue.Enqueue(&newCoord)
				}

				cache[coord.row][coord.col] = cacheQueue
			}
		}

		queue = newQueue
		moves++
	}

	// count 'O' in matrix
	count := 0
	for r := range matrix {
		for _, b := range matrix[r] {
			if b == 'O' {
				count++
			}
		}
	}

	return count
}

func BFS3(matrix [][]byte, root Coord, maxSteps int) int {
	height := len(matrix)
	width := len(matrix[0])
	queue := utils.NewQueue[*Coord]()

	matrix[root.row][root.col] = 'O'
	queue.Enqueue(&root)

	moves := 0
	for moves < maxSteps {
		//fmt.Println(moves + 1)
		newQueue := utils.NewQueue[*Coord]()
		coordInQueue := map[Coord]bool{}

		for queue.HasElement() {
			coord := queue.Dequeue()
			matrix[coord.row][coord.col] = '.'

			if coord.row > 0 && matrix[coord.row-1][coord.col] != '#' {
				newCoord := Coord{coord.row - 1, coord.col}
				if !coordInQueue[newCoord] {
					newQueue.Enqueue(&newCoord)
					matrix[newCoord.row][newCoord.col] = 'O'
					coordInQueue[newCoord] = true
				}
			}

			if coord.row < height-1 && matrix[coord.row+1][coord.col] != '#' {
				newCoord := Coord{coord.row + 1, coord.col}
				if !coordInQueue[newCoord] {
					newQueue.Enqueue(&newCoord)
					matrix[newCoord.row][newCoord.col] = 'O'
					coordInQueue[newCoord] = true
				}
			}

			if coord.col > 0 && matrix[coord.row][coord.col-1] != '#' {
				newCoord := Coord{coord.row, coord.col - 1}
				if !coordInQueue[newCoord] {
					newQueue.Enqueue(&newCoord)
					matrix[newCoord.row][newCoord.col] = 'O'
					coordInQueue[newCoord] = true
				}
			}

			if coord.col < width-1 && matrix[coord.row][coord.col+1] != '#' {
				newCoord := Coord{coord.row, coord.col + 1}
				if !coordInQueue[newCoord] {
					newQueue.Enqueue(&newCoord)
					matrix[newCoord.row][newCoord.col] = 'O'
					coordInQueue[newCoord] = true
				}
			}
		}

		queue = newQueue
		moves++
	}

	// count 'O' in matrix
	count := 0
	for r := range matrix {
		for _, b := range matrix[r] {
			if b == 'O' {
				count++
			}
		}
	}

	return count
}

func Part1(lines []string, maxSteps int) int {
	matrix := make([][]byte, len(lines))

	var root Coord
	for r, line := range lines {
		matrix[r] = make([]byte, len(line))
		for c, b := range line {
			matrix[r][c] = byte(b)
			if b == 'S' {
				root.row = r
				root.col = c
			}
		}
	}

	result := BFS3(matrix, root, maxSteps)

	return result
}

// Solved thanks to the post ox aexl on the reddit thread
// https://www.reddit.com/r/adventofcode/comments/18nevo3/2023_day_21_solutions/
// It implies the creation of a specific input to have a grid large enough to simulate it is infinite
// for the 3 iterations.
/*
We need to figure out how many garden plots can be reached after 26501365 steps.
Note that 26501365 = 202300 * 131 + 65, where 131 is the side length of the input grid.
Store how many garden plots can be reached after 65, 65 + 131 and 65 + 2 * 131 steps, let's call these numbers r₁, r₂ and r₃.
Now it turns out that the number of garden plots that can be reached after x * 131 + 65 steps is a quadratic function p(x) = ax² + bx + c.
We know that
r₁ = p(0) = c
r₂ = p(1) = a + b + c
r₃ = p(2) = 4a + 2b + c
Solving this linear system of equations for the unknowns a, b and c gives
a = (r₃ + r₁ - 2r₂) / 2
b = (4r₂ - 3r₁ - r₃) / 2
c = r₁
Finally, we just need to evaluate the polynomial p at 202300
*/
func Part2(lines []string) int64 {

	nbSteps := 26501365
	gridSize := 131 // from puzzle input
	modulo := nbSteps % gridSize
	leftOver := nbSteps / gridSize

	r1 := Part1(lines, modulo)
	r2 := Part1(lines, modulo+gridSize)
	r3 := Part1(lines, modulo+2*gridSize)

	a := (r3 + r1 - 2*r2) / 2
	b := (4*r2 - 3*r1 - r3) / 2
	c := r1

	p := func(x int) int64 { return int64(a*x*x + b*x + c) }

	return p(leftOver)
}