package main

import (
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/RyanCarrier/dijkstra"
)

func main() {
	inputBytes, _ := os.ReadFile("input.txt")
	input := strings.Split(string(inputBytes), "\n")

	log.Println("Puzzle 1:", puzzle1(input))
	log.Println("Puzzle 2:", puzzle2(input))
}

func puzzle1(input []string) int {
	return lowestRisk(parseInput(input))
}

func puzzle2(input []string) int {
	grid := parseInput(input)
	// Enlarge the grid
	newGrid := [][]int{}
	// Enlarge each row
	for _, row := range grid {
		newRow := []int{}
		for j := 0; j < 5; j++ {
			for _, cell := range row {
				v := cell + j
				if v > 9 {
					v %= 9
				}
				newRow = append(newRow, v)
			}
		}
		newGrid = append(newGrid, newRow)
	}
	// Add more rows
	newNewGrid := [][]int{}
	for i := 0; i < 5; i++ {
		for _, row := range newGrid {
			newRow := []int{}
			for _, cell := range row {
				v := cell + i
				if v > 9 {
					v %= 9
				}
				newRow = append(newRow, v)
			}
			newNewGrid = append(newNewGrid, newRow)
		}
	}
	// Return the lowest risk
	return lowestRisk(newNewGrid)
}

func parseInput(input []string) (grid [][]int) {
	for _, line := range input {
		gridLine := []int{}
		for _, char := range line {
			gridLine = append(gridLine, int(char)-int('0'))
		}
		grid = append(grid, gridLine)
	}
	return
}

func lowestRisk(grid [][]int) int {
	graph := dijkstra.NewGraph()
	// Add all vertices
	for y, row := range grid {
		for x := range row {
			graph.AddMappedVertex(fmt.Sprintf("%d-%d", x, y))
		}
	}
	// Add all edges
	for y, row := range grid {
		for x, cell := range row {
			m := fmt.Sprintf("%d-%d", x, y)
			if y > 0 {
				graph.AddMappedArc(fmt.Sprintf("%d-%d", x, y-1), m, int64(cell))
			}
			if y < len(grid)-1 {
				graph.AddMappedArc(fmt.Sprintf("%d-%d", x, y+1), m, int64(cell))
			}
			if x > 0 {
				graph.AddMappedArc(fmt.Sprintf("%d-%d", x-1, y), m, int64(cell))
			}
			if x < len(row)-1 {
				graph.AddMappedArc(fmt.Sprintf("%d-%d", x+1, y), m, int64(cell))
			}
		}
	}
	// Get shortest path
	best, err := graph.Shortest(graph.AddMappedVertex("0-0"), graph.AddMappedVertex(fmt.Sprintf("%d-%d", len(grid[0])-1, len(grid)-1)))
	if err != nil {
		log.Fatal(err)
	}
	// Return the distance
	return int(best.Distance)
}