package main

import (
	"fmt"
	"math"
	"strings"
)

type Position struct {
	x, y int
}

type Movement struct {
	from Position
	to   Position
}

type Keypad struct {
	keys      map[byte]Position
	forbidden Position
}

var numpad = Keypad{
	keys: map[byte]Position{
		'7': {0, 0},
		'8': {1, 0},
		'9': {2, 0},
		'4': {0, 1},
		'5': {1, 1},
		'6': {2, 1},
		'1': {0, 2},
		'2': {1, 2},
		'3': {2, 2},
		'0': {1, 3},
		'A': {2, 3},
	},
	forbidden: Position{0, 3},
}

var dirpad = Keypad{
	keys: map[byte]Position{
		'^': {1, 0},
		'A': {2, 0},
		'<': {0, 1},
		'v': {1, 1},
		'>': {2, 1},
	},
	forbidden: Position{0, 0},
}

type DpKey struct {
	code  string
	depth int
}

func main() {
	fmt.Println(part1("029A"), 68)
	fmt.Println(part1("980A"), 60)
	fmt.Println(part1("179A"), 68)
	fmt.Println(part1("456A"), 64)
	fmt.Println(part1("379A"), 64)
	fmt.Println("")
	fmt.Println(137870)
	fmt.Println(part1("805A")*805 + part1("170A")*170 + part1("129A")*129 + part1("283A")*283 + part1("540A")*540)
	fmt.Println("")
	fmt.Println(part2("029A"))
	fmt.Println(part2("980A"))
	fmt.Println(part2("179A"))
	fmt.Println(part2("456A"))
	fmt.Println(part2("379A"))
	fmt.Println("")
	fmt.Println(170279148659464)
	fmt.Println(part2("805A")*805 + part2("170A")*170 + part2("129A")*129 + part2("283A")*283 + part2("540A")*540)
}

func part1(code string) int {
	memo := make(map[DpKey]int)
	return typeCode(memo, code, []Keypad{numpad, dirpad, dirpad}, 0)
}

func part2(code string) int {
	memo := make(map[DpKey]int)
	return typeCode(memo, code, []Keypad{
		numpad,
		dirpad, dirpad, dirpad, dirpad, dirpad,
		dirpad, dirpad, dirpad, dirpad, dirpad,
		dirpad, dirpad, dirpad, dirpad, dirpad,
		dirpad, dirpad, dirpad, dirpad, dirpad,
		dirpad, dirpad, dirpad, dirpad, dirpad,
	}, 0)
}

func combine(prefix []string, suffix []string) (res []string) {
	for _, p := range prefix {
		for _, s := range suffix {
			res = append(res, p+s)
		}
	}
	return
}

func typeCode(memo map[DpKey]int, code string, keypads []Keypad, depth int) int {
	subCode := strings.Split(code, "A")
	subCode = subCode[:len(subCode)-1]
	cost := 0
	for _, subCode := range subCode {
		cost += typeSubCode(memo, subCode+"A", keypads, depth)
	}
	return cost
}

func typeSubCode(memo map[DpKey]int, code string, keypads []Keypad, depth int) int {
	if cost, isIn := memo[DpKey{code, depth}]; isIn {
		return cost
	}

	codeSeqs := []string{}
	{
		pos := keypads[depth].keys['A']
		seqs := []string{""}
		for i := range code {
			nextPos := keypads[depth].keys[code[i]]
			childSeqs := buildSeqs(keypads[depth], pos, nextPos)
			seqs = combine(seqs, childSeqs)
			pos = nextPos
		}
		codeSeqs = append(codeSeqs, seqs...)
	}

	minCost := math.MaxInt
	if depth < len(keypads)-1 {
		for _, seq := range codeSeqs {
			minCost = min(typeCode(memo, seq, keypads, depth+1), minCost)
		}
	} else {
		for _, seq := range codeSeqs {
			minCost = min(len(seq), minCost)
		}
	}

	memo[DpKey{code, depth}] = minCost
	return minCost
}

func shortestSeq(seqs []string) int {
	minLength := math.MaxInt
	for _, s := range seqs {
		minLength = min(minLength, len(s))
	}
	return minLength
}

func buildSeqs(kp Keypad, from, to Position) (seqs []string) {
	type Step struct {
		p   Position
		seq string
	}
	steps := []Step{{from, ""}}
	for len(steps) > 0 {
		nextSteps := []Step{}
		for _, s := range steps {
			if s.p == to {
				seqs = append(seqs, s.seq+"A")
				continue
			}
			if s.p.x < to.x && !(s.p.x+1 == kp.forbidden.x && s.p.y == kp.forbidden.y) {
				nextSteps = append(nextSteps, Step{Position{s.p.x + 1, s.p.y}, s.seq + ">"})
			} else if s.p.x > to.x && !(s.p.x-1 == kp.forbidden.x && s.p.y == kp.forbidden.y) {
				nextSteps = append(nextSteps, Step{Position{s.p.x - 1, s.p.y}, s.seq + "<"})
			}
			if s.p.y < to.y && !(s.p.y+1 == kp.forbidden.y && s.p.x == kp.forbidden.x) {
				nextSteps = append(nextSteps, Step{Position{s.p.x, s.p.y + 1}, s.seq + "v"})
			} else if s.p.y > to.y && !(s.p.y-1 == kp.forbidden.y && s.p.x == kp.forbidden.x) {
				nextSteps = append(nextSteps, Step{Position{s.p.x, s.p.y - 1}, s.seq + "^"})
			}
		}
		steps = nextSteps
	}
	if len(seqs) == 0 {
		seqs = []string{"A"}
	}
	return
}