package main

import (
	"bufio"
	"container/heap"
	"fmt"
	"math"
	"os"
	"strings"
)

func main() {
	fmt.Println(7036, 45)
	fmt.Println(part1And2(readData("example.txt")))
	fmt.Println(11048, 64)
	fmt.Println(part1And2(readData("example2.txt")))
	fmt.Println(89460, 504)
	fmt.Println(part1And2(readData("data.txt")))
}

type Position struct {
	x, y int
}

type PositionDir struct {
	x, y int
	dir  int
}

type Candidate struct {
	x, y    int
	dir     int
	cost    int
	visited []PositionDir
}

type CandidateQueue struct {
	elements []Candidate
}

func (queue *CandidateQueue) Len() int {
	return len(queue.elements)
}

func (queue *CandidateQueue) Less(i, j int) bool {
	return queue.elements[i].cost < queue.elements[j].cost
}

func (queue *CandidateQueue) Swap(i, j int) {
	queue.elements[i], queue.elements[j] = queue.elements[j], queue.elements[i]
}

func (queue *CandidateQueue) Push(x any) {
	queue.elements = append(queue.elements, x.(Candidate))
}

func (queue *CandidateQueue) AddElement(candidate Candidate) {
	heap.Push(queue, candidate)
}

func (queue *CandidateQueue) Pop() any {
	if len(queue.elements) == 0 {
		return nil
	} else {
		elmt := queue.elements[len(queue.elements)-1]
		queue.elements = queue.elements[:len(queue.elements)-1]
		return elmt
	}
}

func cloneVisited(v []PositionDir) []PositionDir {
	clone := make([]PositionDir, len(v))
	for i, p := range v {
		clone[i] = p
	}
	return clone
}

var directions = [4][2]int{{1, 0}, {0, 1}, {-1, 0}, {0, -1}}

// Originally I had implemented A* for part1. But because of part2 I modified it to be
// dumber. I wish we had full requirements from the start so I wouldn't waste time...

func part1And2(maze [][]byte) (scorePart1, scorePart2 int) {
	endX, endY := len(maze[0])-2, 1
	startX, startY := 1, len(maze)-2

	queue := new(CandidateQueue)
	heap.Push(queue, Candidate{startX, startY, 0, 0, []PositionDir{{startX, startY, 0}}})
	minCost := math.MaxInt
	visitedCost := make(map[PositionDir]int)
	onOptimalPath := make(map[Position]bool)

	for queue.Len() > 0 {
		candidate := heap.Pop(queue).(Candidate)
		posDir := PositionDir{candidate.x, candidate.y, candidate.dir}

		alreadyVisitedCost, alreadyVisited := visitedCost[posDir]
		if alreadyVisited && alreadyVisitedCost < candidate.cost {
			continue
		}
		visitedCost[posDir] = candidate.cost

		tryDirection := func(dir int, cost int) {
			fwdX := candidate.x + directions[dir][0]
			fwdY := candidate.y + directions[dir][1]
			realCostFwd := candidate.cost + cost
			if maze[fwdY][fwdX] != '#' && realCostFwd <= minCost {
				if fwdX == endX && fwdY == endY && realCostFwd <= minCost {
					for _, p := range candidate.visited {
						onOptimalPath[Position{p.x, p.y}] = true
					}
					minCost = realCostFwd
				} else {
					visited := append(cloneVisited(candidate.visited), PositionDir{fwdX, fwdY, dir})
					c := Candidate{fwdX, fwdY, dir, realCostFwd, visited}
					heap.Push(queue, c)
				}
			}
		}

		tryDirection(candidate.dir, 1)
		tryDirection((candidate.dir+1)%4, 1001)
		tryDirection((candidate.dir+3)%4, 1001)
	}

	return minCost, len(onOptimalPath) + 1
}

func readData(fileName string) (maze [][]byte) {
	fp, _ := os.Open(fileName)
	scanner := bufio.NewScanner(fp)

	for scanner.Scan() {
		maze = append(maze, []byte(strings.TrimSpace(scanner.Text())))
	}
	return
}