package pacificatlanticwaterflow

type Position struct{ X, Y int }

func (p Position) Neighbors() []Position {
	return []Position{
		{p.X + 1, p.Y},
		{p.X - 1, p.Y},
		{p.X, p.Y + 1},
		{p.X, p.Y - 1},
	}
}

func (p Position) ToList() []int {
	return []int{p.Y, p.X}
}

type Board [][]int

func (b Board) Get(p Position) int {
	if b.OutOfBounds(p) {
		return -1
	}

	return b[p.Y][p.X]
}

func (b Board) Height() int { return len(b) }
func (b Board) Width() int  { return len(b[0]) }

func (b Board) OutOfBounds(p Position) bool {
	return p.Y < 0 || p.Y >= b.Height() || p.X < 0 || p.X >= b.Width()
}

type Flow struct{ Pacific, Atlantic bool }

func (f *Flow) Merge(o Flow) {
	f.Atlantic = f.Atlantic || o.Atlantic
	f.Pacific = f.Pacific || o.Pacific
}

func (f Flow) Both() bool {
	return f.Pacific && f.Atlantic
}

func (b Board) GetFlow(p Position) Flow {
	return Flow{
		Atlantic: p.Y >= b.Height() || p.X >= b.Width(),
		Pacific:  p.Y < 0 || p.X < 0,
	}
}

func pacificAtlanticHelper(board Board, position Position, flows map[Position]Flow) Flow {
	if flow, ok := flows[position]; ok {
		return flow
	}

	if board.OutOfBounds(position) {
		return board.GetFlow(position)
	}

	flow := Flow{}

	var top Position
	plateau := map[Position]bool{}
	unvisited := []Position{position}
	for len(unvisited) > 0 {
		top, unvisited = unvisited[len(unvisited)-1], unvisited[:len(unvisited)-1]

		if board.Get(top) < board.Get(position) {
			flow.Merge(pacificAtlanticHelper(board, top, flows))
		} else if board.Get(top) == board.Get(position) {
			if _, ok := plateau[top]; !ok {
				plateau[top] = true
				unvisited = append(unvisited, top.Neighbors()...)
			}
		}
	}

	for mem := range plateau {
		flows[mem] = flow
	}

	return flow
}

func PacificAtlantic(heights [][]int) [][]int {
	flows := map[Position]Flow{}
	board := Board(heights)

	for y := range board.Height() {
		for x := range board.Width() {
			pacificAtlanticHelper(board, Position{x, y}, flows)
		}
	}

	result := [][]int{}
	for position, flow := range flows {
		if flow.Both() {
			result = append(result, position.ToList())
		}
	}

	return result
}