package day17

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

type NodeID int
type Distance int

type Node struct {
	nodeID    NodeID
	dist      Distance
	direction rune
}

type Neighbor struct {
	nodeID    NodeID
	dist      Distance
	direction rune
}

type Step struct {
	nodeID     NodeID
	dist       Distance
	direction  rune
	stepsInDir int
}

func makeStepKey(s *Step) string {
	return fmt.Sprintf("%v", s)
}

func findValidNeighborSteps(neighbors []Neighbor, step *Step, visited map[string]bool, min int, max int) []*Step {
	result := []*Step{}

	for _, neighbor := range neighbors {
		/*
			if known[neighbor.nodeID] {
				continue
			}
		*/

		if neighbor.direction != step.direction && step.stepsInDir < min {
			continue
		}

		if neighbor.direction == step.direction && step.stepsInDir == max {
			continue
		}

		if neighbor.direction == 'U' && step.direction == 'D' || neighbor.direction == 'D' && step.direction == 'U' ||
			neighbor.direction == 'L' && step.direction == 'R' || neighbor.direction == 'R' && step.direction == 'L' {
			continue
		}

		countSteps := 0
		if neighbor.direction == step.direction {
			countSteps = step.stepsInDir + 1
		} else {
			countSteps = 1
		}

		newStep := &Step{
			nodeID:     neighbor.nodeID,
			dist:       neighbor.dist,
			direction:  neighbor.direction,
			stepsInDir: countSteps,
		}

		newStepKey := makeStepKey(newStep)
		if visited[newStepKey] {
			continue
		}
		result = append(result, newStep)
	}

	return result
}

func dijkstraShortestPath(adjacencyList [][]Neighbor, start, end NodeID, min int, max int) (map[string]*Step, map[string]Distance) {
	vectorSize := len(adjacencyList)
	prev := make(map[string]*Step, vectorSize)
	dist := make(map[string]Distance, vectorSize)
	queue := make(map[string]*Step, vectorSize)
	visited := make(map[string]bool, vectorSize)

	step0R := Step{
		nodeID:     start,
		dist:       0,
		direction:  'R',
		stepsInDir: 0,
	}
	step0D := Step{
		nodeID:     start,
		dist:       0,
		direction:  'D',
		stepsInDir: 0,
	}

	dist[makeStepKey(&step0R)] = 0
	dist[makeStepKey(&step0D)] = 0

	queue[makeStepKey(&step0R)] = &step0R
	queue[makeStepKey(&step0D)] = &step0D

	for len(queue) > 0 {
		var key string
		var step *Step
		minDist := Distance(math.MaxInt32)

		for k, s := range queue {

			currDist, found := dist[k]
			if found && currDist < minDist {
				minDist = currDist
				key = k
				step = s
			}
		}

		if step.nodeID == end {
			break
		}

		delete(queue, key)
		visited[key] = true

		neighborSteps := findValidNeighborSteps(adjacencyList[step.nodeID], step, visited, min, max)

		for _, neighborStep := range neighborSteps {
			queue[makeStepKey(neighborStep)] = neighborStep

			alt := minDist + neighborStep.dist

			neighborKey := makeStepKey(neighborStep)
			neighborDist, found := dist[neighborKey]
			if !found || alt < neighborDist {
				dist[neighborKey] = alt
				prev[neighborKey] = step
			}
		}
	}

	return prev, dist
}

func MaxHeat(lines []string, min int, max int) int {
	height := len(lines)
	width := len(lines[0])
	coordToID := func(row, col int) NodeID {
		return NodeID(row*width + col)
	}

	adjacencyList := make([][]Neighbor, height*width)

	for row := range lines {
		for col := range lines[row] {

			nodeID := coordToID(row, col)

			neighbors := []Neighbor{}

			if row > 0 {
				dist := lines[row-1][col] - '0'
				neighbors = append(neighbors, Neighbor{
					nodeID:    coordToID(row-1, col),
					dist:      Distance(dist),
					direction: 'U',
				})
			}

			if row < height-1 {
				dist := lines[row+1][col] - '0'
				neighbors = append(neighbors, Neighbor{
					nodeID:    coordToID(row+1, col),
					dist:      Distance(dist),
					direction: 'D',
				})
			}

			if col > 0 {
				dist := lines[row][col-1] - '0'
				neighbors = append(neighbors, Neighbor{
					nodeID:    coordToID(row, col-1),
					dist:      Distance(dist),
					direction: 'L',
				})
			}

			if col < width-1 {
				dist := lines[row][col+1] - '0'
				neighbors = append(neighbors, Neighbor{
					nodeID:    coordToID(row, col+1),
					dist:      Distance(dist),
					direction: 'R',
				})
			}

			adjacencyList[nodeID] = neighbors
		}
	}

	start := coordToID(0, 0)
	end := coordToID(height-1, width-1)
	prevs, distances := dijkstraShortestPath(adjacencyList, start, end, min, max)

	minDist := Distance(math.MaxInt32)
	minKey := ""
	for key, dist := range distances {
		if strings.HasPrefix(key, fmt.Sprintf("&{%d", end)) {
			fmt.Println(key, dist)
			if dist < minDist {
				minDist = dist
				minKey = key
			}
		}
	}

	fmt.Print(minKey, "-")
	step, found := prevs[minKey]
	for found {
		fmt.Printf("(%d,%d) %s - ", int(step.nodeID)/width, int(step.nodeID)%width, string(step.direction))
		step, found = prevs[makeStepKey(step)]
	}

	/*
		nodeID := end
		for nodeID != start {
			fmt.Printf("(%d,%d) - ", int(nodeID)/width, int(nodeID)%width)
			nodeID = prevs[nodeID]
		}
	*/
	return int(minDist) //int(distances[end])
}

var adjacencyList [][]Node
var visited []bool
var currentPath []Node
var simplePath [][]Node
var distances []Distance
var vectorSize int
var currMin Distance
var minimums map[string]Distance

func key(node1, node2 NodeID) string { return fmt.Sprintf("%d-%d", node1, node2) }

func findValidNeighbors(neighbors []Node, directionsTaken string) []Node {
	result := []Node{}

	for _, neighbor := range neighbors {
		destDirection := directionsTaken + string(neighbor.direction)
		lastFour := ""
		if len(destDirection) >= 4 {
			lastFour = destDirection[len(destDirection)-4:]
		}

		if lastFour == "UUUU" || lastFour == "DDDD" || lastFour == "LLLL" || lastFour == "RRRR" {
			continue
		}

		result = append(result, neighbor)
	}

	return result
}

func DFS(current Node, goalID NodeID, previousDist Distance, directionsTaken string) {

	if previousDist+current.dist > currMin {
		return
	}
	/*
		k := key(current.nodeID, goalID)
		minDist, found := minimums[k]
		if found && minDist < previousDist+current.dist {
			return
		}
	*/
	if len(simplePath) > 1000 {
		log.Fatalf("loop 1 ?")
	}

	if len(directionsTaken) > vectorSize {
		log.Fatalf("loop 2 ?")
	}

	if visited[current.nodeID] {
		return
	}

	visited[current.nodeID] = true
	currentPath = append(currentPath, current)

	if current.nodeID == goalID {
		distance := Distance(0)
		for i := len(currentPath) - 1; i >= 0; i-- {
			distance += currentPath[i].dist
			if currentPath[i].nodeID != goalID {
				k := key(currentPath[i].nodeID, goalID)
				minimums[k] = distance
			}
		}
		//distance := utils.Reduce(currentPath, 0, func(acc Distance, curr Node) Distance { return acc + curr.dist })
		if distance < currMin {
			distances = append(distances, distance)
			duplicate := make([]Node, len(currentPath))
			copy(duplicate, currentPath)
			simplePath = append(simplePath, duplicate)
			currMin = distance

		}
		visited[current.nodeID] = false
		currentPath = currentPath[:len(currentPath)-1]
		return
	}

	neighbors := findValidNeighbors(adjacencyList[current.nodeID], directionsTaken)

	for _, neighbor := range neighbors {
		DFS(neighbor, goalID, previousDist+current.dist, directionsTaken+string(neighbor.direction))
	}

	visited[current.nodeID] = false
	currentPath = currentPath[:len(currentPath)-1]
}

func MaxHeat2(lines []string) int {
	height := len(lines)
	width := len(lines[0])

	coordToID := func(row, col int) NodeID { return NodeID(row*width + col) }

	vectorSize = height * width
	visited = make([]bool, vectorSize)
	directionsTaken := ""
	adjacencyList = make([][]Node, vectorSize)
	currMin = Distance(math.MaxInt32)

	for row := range lines {
		for col := range lines[row] {

			nodeID := coordToID(row, col)

			neighbors := []Node{}

			if row > 0 {
				dist := lines[row-1][col] - '0'
				neighbors = append(neighbors, Node{
					nodeID:    coordToID(row-1, col),
					dist:      Distance(dist),
					direction: 'U',
				})
			}

			if row < height-1 {
				dist := lines[row+1][col] - '0'
				neighbors = append(neighbors, Node{
					nodeID:    coordToID(row+1, col),
					dist:      Distance(dist),
					direction: 'D',
				})
			}

			if col > 0 {
				dist := lines[row][col-1] - '0'
				neighbors = append(neighbors, Node{
					nodeID:    coordToID(row, col-1),
					dist:      Distance(dist),
					direction: 'L',
				})
			}

			if col < width-1 {
				dist := lines[row][col+1] - '0'
				neighbors = append(neighbors, Node{
					nodeID:    coordToID(row, col+1),
					dist:      Distance(dist),
					direction: 'R',
				})
			}

			adjacencyList[nodeID] = neighbors
		}
	}

	start := Node{
		nodeID:    coordToID(0, 0),
		dist:      0,
		direction: ' ',
	}
	goalID := coordToID(height-1, width-1)
	currentPath = []Node{}
	simplePath = [][]Node{}
	distances = []Distance{}
	minimums = map[string]Distance{}

	DFS(start, goalID, 0, directionsTaken)

	minDist := Distance(math.MaxInt32)
	idx := -1
	for i, v := range distances {
		if v < minDist {
			idx = i
			minDist = v
		}
	}

	for _, node := range simplePath[idx] {
		fmt.Printf("(%d,%d) - ", int(node.nodeID)/width, int(node.nodeID)%width)
	}

	return int(minDist)
}