package day21

import (
	"adventofcode2024/utils"
	"adventofcode2024/utils/dijkstra"
	"adventofcode2024/utils/memo"
	"fmt"
	"math"
	"strings"
)

type Pair struct {
	s dijkstra.Point
	d dijkstra.Point
}

type Path struct {
	path     []dijkstra.Point
	path_str string
	cost     int
}

type State struct {
	code         string
	depth        int
	isNumericPad bool
}

var numericPad = map[dijkstra.Point]string{
	{X: 0, Y: 0}: "7", {X: 1, Y: 0}: "8", {X: 2, Y: 0}: "9",
	{X: 0, Y: 1}: "4", {X: 1, Y: 1}: "5", {X: 2, Y: 1}: "6",
	{X: 0, Y: 2}: "1", {X: 1, Y: 2}: "2", {X: 2, Y: 2}: "3",
	{X: 1, Y: 3}: "0", {X: 2, Y: 3}: "A",
}

var directionPad = map[dijkstra.Point]string{
	{X: 1, Y: 0}: "^", {X: 2, Y: 0}: "A",
	{X: 0, Y: 1}: "<", {X: 1, Y: 1}: "v", {X: 2, Y: 1}: ">",
}

var directions = []dijkstra.Point{{X: 0, Y: 1}, {X: 0, Y: -1}, {X: 1, Y: 0}, {X: -1, Y: 0}}

var lookup_np = get_paths(numericPad)
var lookup_dp = get_paths(directionPad)

var m *memo.Memo

func init() {
	m = memo.New(get_shortest_path)
}

func Part1(input string) int {
	codes := strings.Split(input, "\n")
	cost := 0
	for _, code := range codes {
		initialState := State{
			code:         code,
			depth:        2,
			isNumericPad: true,
		}
		result := m.Get(initialState).(int)
		cost += result * utils.MustAtoi(code[:len(code)-1])
	}
	return cost
}

func Part2(input string) int {
	codes := strings.Split(input, "\n")
	cost := 0
	for _, code := range codes {
		initialState := State{
			code:         code,
			depth:        25,
			isNumericPad: true,
		}
		result := m.Get(initialState).(int)
		cost += result * utils.MustAtoi(code[:len(code)-1])
	}
	return cost
}

func get_paths(in map[dijkstra.Point]string) map[Pair][]Path {
	lookup := make(map[Pair][]Path)
	data := dijkstra.InputGraph{}
	for p := range in {
		for _, dir := range directions {
			np := dijkstra.Point{X: p.X + dir.X, Y: p.Y + dir.Y}
			if in[np] != "" {
				data.Graph = append(data.Graph, dijkstra.InputData{
					Source:      p,
					Destination: np,
					Weight:      1,
				})
			}
		}
	}
	graph := dijkstra.CreateGraph(data)
	for s := range in {
		for d := range in {
			paths, cost := dijkstra.GetAllShortestPaths(s, d, graph)

			for _, path := range paths {
				path_str := ""
				for x := 1; x < len(path); x++ {
					f := path[x-1]
					t := path[x]

					switch {
					case f.X == t.X && f.Y < t.Y:
						// Moving verticallY down
						path_str += "v"
					case f.X == t.X && f.Y > t.Y:
						// Moving verticallY up
						path_str += "^"
					case f.Y == t.Y && f.X < t.X:
						// Moving horizontallY right
						path_str += ">"
					case f.Y == t.Y && f.X > t.X:
						// Moving horizontally left
						path_str += "<"
					default:
						// No matching case or invalid input
						fmt.Println("Invalid direction")
					}

				}
				lookup[Pair{s, d}] = append(lookup[Pair{s, d}], Path{path, path_str, cost})
			}

		}
	}
	return lookup
}

func find_point(in string, lookup map[dijkstra.Point]string) dijkstra.Point {
	for k, v := range lookup {
		if v == in {
			return k
		}
	}
	return dijkstra.Point{}
}

func get_shortest_path(key interface{}) interface{} {
	state := key.(State)

	if state.depth < 0 {
		return len(state.code)
	}

	// (Optional) If you need to modify the state, work on a copy.
	newState := state
	newState.code = "A" + state.code

	cost := 0
	for x := 0; x<len(newState.code)-1; x++ {
		pps := shortest_paths(string(newState.code[x]), string(newState.code[x+1]), newState.isNumericPad)
		min := math.MaxInt // Start with a very high number not just 65535
		for _, pp := range pps {
			nextState := State{
				code:         pp.path_str + "A",
				depth:        newState.depth - 1,
				isNumericPad: false,
			}
			mm := m.Get(nextState).(int)
			if mm < min {
				min = mm
			}
		}
		cost += min
	}
	return cost
}

func shortest_paths(s, d string, isNumericPad bool) []Path {
	if isNumericPad {
		return lookup_np[Pair{s: find_point(s, numericPad), d: find_point(d, numericPad)}]
	} else {
		return lookup_dp[Pair{s: find_point(s, directionPad), d: find_point(d, directionPad)}]
	}
}