package main

import (
	"fmt"

	h "git.bullercodeworks.com/brian/adventofcode/helpers"
)

func main() {
	inp := h.StdinToStringSlice()
	part1(inp)
	fmt.Println()
	part2(inp)
}

func part1(input []string) {
	m := h.StringSliceToCoordByteMap(input)

	fmt.Println("# Part 1")
	fmt.Println(solve(&m, 0, 3))
}

// 1188 is too low
func part2(input []string) {
	m := h.StringSliceToCoordByteMap(input)

	fmt.Println("# Part 2")
	fmt.Println(solve(&m, 4, 10))
}

func solve(m *h.CoordByteMap, minStraight, maxStraight int) int {
	q := NewPriorityQueue[heatState](func(a, b heatState) int {
		return a.heatLoss - b.heatLoss
	})
	q.Push(heatState{
		state: state{
			pos:      &h.Coordinate{X: m.TLX + 1, Y: m.TLY},
			dir:      dirE,
			dirCount: 1,
		},
	})
	q.Push(heatState{
		state: state{
			pos:      &h.Coordinate{X: m.TLX, Y: m.TLY + 1},
			dir:      dirS,
			dirCount: 1,
		},
	})
	end := h.Coordinate{X: m.BRX, Y: m.BRY}
	visitedCoords := make(map[h.Coordinate]bool)
	visited := make(map[string]int)

	for q.IsNotEmpty() {
		chk := q.Pop()
		if !m.ContainsCoord(*chk.state.pos) {
			continue
		}
		heat := int(m.Get(*chk.state.pos)-'0') + chk.heatLoss
		if chk.state.pos.Equals(end) && chk.dirCount >= minStraight {
			return heat
		}
		if visHeat, ok := visited[chk.state.String()]; ok {
			if visHeat <= heat {
				continue
			}
		}
		visitedCoords[*chk.state.pos] = true
		visited[chk.state.String()] = heat
		if chk.dirCount >= minStraight {
			// Add Left Turn
			left := chk.dir.LeftTurn().Get(chk.pos)
			if m.ContainsCoord(*left) {
				q.Push(heatState{
					state: state{
						pos:      left,
						dir:      chk.dir.LeftTurn(),
						dirCount: 1,
					},
					heatLoss: heat,
				})
			}
			// Add Right Turn
			right := chk.dir.RightTurn().Get(chk.pos)
			if m.ContainsCoord(*right) {
				q.Push(heatState{
					state: state{
						pos:      right,
						dir:      chk.dir.RightTurn(),
						dirCount: 1,
					},
					heatLoss: heat,
				})
			}
		}

		if chk.dirCount < maxStraight {
			// Add Straight
			straight := chk.dir.Get(chk.pos)
			if m.ContainsCoord(*straight) {
				q.Push(heatState{
					state: state{
						pos:      straight,
						dir:      chk.dir,
						dirCount: chk.dirCount + 1,
					},
					heatLoss: heat,
				})
			}
		}
	}
	return h.MAX_INT
}

type state struct {
	pos      *h.Coordinate
	dir      direction
	dirCount int
}

func (s state) String() string {
	return fmt.Sprintf("%s-%s-%d", s.pos, s.dir, s.dirCount)
}

type heatState struct {
	state
	heatLoss int
}

type PriorityQueue[T comparable] struct {
	items      []T
	comparator func(a, b T) int
}

func NewPriorityQueue[T comparable](comparator func(a, b T) int) PriorityQueue[T] {
	return PriorityQueue[T]{comparator: comparator}
}
func (pq *PriorityQueue[T]) Push(item T) {
	pq.items = append(pq.items, item)
	pq.heapifyUp(len(pq.items) - 1)
}
func (pq *PriorityQueue[T]) Pop() T {
	top := pq.items[0]
	lastIndex := len(pq.items) - 1
	pq.items[0], pq.items[lastIndex] = pq.items[lastIndex], pq.items[0]
	pq.items = pq.items[:lastIndex]
	pq.heapifyDown(0)
	return top
}
func (pq *PriorityQueue[T]) Peek() T {
	return pq.items[0]
}
func (pq *PriorityQueue[T]) Len() int {
	return len(pq.items)
}
func (pq *PriorityQueue[T]) IsEmpty() bool {
	return len(pq.items) == 0
}
func (pq *PriorityQueue[T]) IsNotEmpty() bool {
	return !pq.IsEmpty()
}
func (pq *PriorityQueue[T]) heapifyUp(index int) {
	for index > 0 {
		parentIndex := (index - 1) / 2
		if pq.comparator(pq.items[index], pq.items[parentIndex]) < 0 {
			pq.items[index], pq.items[parentIndex] = pq.items[parentIndex], pq.items[index]
			index = parentIndex
		} else {
			break
		}
	}
}
func (pq *PriorityQueue[T]) heapifyDown(index int) {
	for {
		leftChild, rightChild, smallest := 2*index+1, 2*index+2, index
		if leftChild < len(pq.items) && pq.comparator(pq.items[leftChild], pq.items[smallest]) < 0 {
			smallest = leftChild
		}
		if rightChild < len(pq.items) && pq.comparator(pq.items[rightChild], pq.items[smallest]) < 0 {
			smallest = rightChild
		}
		if smallest != index {
			pq.items[index], pq.items[smallest] = pq.items[smallest], pq.items[index]
			index = smallest
		} else {
			break
		}
	}
}

type direction h.Coordinate

func (d direction) Get(from *h.Coordinate) *h.Coordinate {
	switch d.String() {
	case "N":
		return from.GetNorthCoord()
	case "E":
		return from.GetEastCoord()
	case "S":
		return from.GetSouthCoord()
	case "W":
		return from.GetWestCoord()
	}
	return from
}
func (d direction) LeftTurn() direction {
	switch d.String() {
	case "N":
		return dirW
	case "E":
		return dirN
	case "S":
		return dirE
	default:
		return dirS
	}
}
func (d direction) RightTurn() direction {
	switch d.String() {
	case "N":
		return dirE
	case "E":
		return dirS
	case "S":
		return dirW
	default:
		return dirN
	}
}
func (d direction) String() string {
	switch {
	case d.X == dirN.X && d.Y == dirN.Y:
		return "N"
	case d.X == dirE.X && d.Y == dirE.Y:
		return "E"
	case d.X == dirS.X && d.Y == dirS.Y:
		return "S"
	default:
		return "W"
	}
}

var (
	dirN = direction{X: 0, Y: -1}
	dirE = direction{X: 1, Y: 0}
	dirS = direction{X: 0, Y: 1}
	dirW = direction{X: -1, Y: 0}
)