"""
Copyright (c) Meta Platforms, Inc. and affiliates.
All rights reserved.

This source code is licensed under the license found in the
LICENSE file in the root directory of this source tree.
"""

from functools import partial

import jax
import jax.numpy as jnp
import numpy as np


@partial(jax.jit, static_argnums=(1,))
def apsp(A, n=None):
	"""
	Compute APSP for adjacency matrix A 
	using Seidel's algorithm.
	"""
	if n is None:
		n = A.shape[0]
	assert(n == A.shape[0]), 'n must equal dim of A.'

	n_steps = int(np.ceil(np.log(n)/np.log(2)))
	A_cache = jnp.zeros((n_steps, n, n), dtype=jnp.uint32)
	steps_to_reduce = jnp.array(1, dtype=jnp.int32)

	def _scan_fwd_step(carry, step): 
		i = step
		A, A_cache, steps_to_reduce = carry
		A_cache = A_cache.at[i].set(A)

		Z = A@A
		B = jnp.logical_or(
				A == 1, 
				Z > 0
			).astype(jnp.uint32) \
			 .at[jnp.diag_indices(n)].set(0)
		A = B

		complete = B.sum() - jnp.diagonal(B).sum() == n*(n-1)
		steps_to_reduce += ~complete

		return (A, A_cache, steps_to_reduce), None

	(B, A_cache, steps_to_reduce), _ = jax.lax.scan(
		_scan_fwd_step,
		(A, A_cache, 1),
		jnp.arange(n_steps),
		length=n_steps
	)

	D = 2*B - A_cache[steps_to_reduce-1]

	def _scan_bkwd_step(carry, step):
		i = step
		(T, A_cache,steps_to_reduce) = carry

		A = A_cache[steps_to_reduce - i - 1]
		X = T@A

		thresh = T*(jnp.tile(A.sum(0, keepdims=True), (n, 1)))
		D = 2*T*(X >= thresh) + (2*T - 1)*(X < thresh)
		T = D*(i < steps_to_reduce) + T*(i >= steps_to_reduce)

		return (T, A_cache, steps_to_reduce), None

	(D, _, _), _ = jax.lax.scan(
		_scan_bkwd_step,
		(D, A_cache, steps_to_reduce),
		jnp.arange(1, n_steps),
		length=n_steps-1
	)

	return D


@jax.jit
def grid_to_graph(grid):
	"""
	Transform a binary grid (True == wall) into a
	graph, where walls are all connected to a default node.
	"""
	h, w = grid.shape
	nodes = grid.flatten()
	n = len(nodes)
	A = jnp.zeros((n,n), dtype=jnp.uint32)
	
	all_idx = jnp.arange(n)
	dum_neighbor_idx = jnp.argmax(~nodes)
	dum_neighbor_mask = jnp.zeros(n, dtype=jnp.bool_)
	dum_neighbor_mask = \
		dum_neighbor_mask.at[dum_neighbor_idx].set(True)

	def _get_neigbors(idx):
		# Return length n boolean mask of neighbors
		# We then vmap this function over all n
		r = idx + 1
		l = idx - 1
		t = idx - w
		b = idx + w

		l_border = jnp.logical_or(
			idx % w == 0,
			nodes[l]
		)
		r_border = jnp.logical_or(
			r % w == 0,
			nodes[r]
		)
		t_border = jnp.logical_or(
			idx // w == 0,
			nodes[t],
		)
		b_border = jnp.logical_or(
			idx // w == h - 1,
			nodes[b]
		)

		l_ignore = jnp.logical_or(
			l_border,
			nodes[idx]
		)
		r_ignore = jnp.logical_or(
			r_border,
			nodes[idx]
		)
		t_ignore = jnp.logical_or(
			t_border,
			nodes[idx]
		)
		b_ignore = jnp.logical_or(
			b_border,
			nodes[idx]
		)

		left = l*(1-l_ignore) + idx*(l_ignore)
		right = r*(1-r_ignore) + idx*(r_ignore)
		top = t*(1-t_ignore) + idx*(t_ignore)
		bottom = b*(1-b_ignore) + idx*(b_ignore)

		neighbor_mask = jnp.zeros(n, dtype=jnp.bool_)
		neighbor_mask = neighbor_mask.at[jnp.array([left, right, top, bottom])].set(True)

		neighbor_mask = (1-nodes[idx])*neighbor_mask + nodes[idx]*dum_neighbor_mask

		neighbor_mask = neighbor_mask.at[idx].set(False)

		return neighbor_mask

	A = jax.vmap(_get_neigbors)(all_idx).astype(dtype=jnp.uint32)
	A = jnp.maximum(A, A.transpose())

	return A
	

NEIGHBOR_OFFSETS = jnp.array([
	[1,0], # right
	[0,1], # bottom
	[-1,0], # left
	[0,-1], # top
	[0,0] # self
], dtype=jnp.int32)


@jax.jit
def component_mask_with_pos(grid, pos_a):
	"""
	Return a mask set to True in all cells that are
	a part of the connected component containing pos_a.
	"""
	h,w = grid.shape
	visited_mask = grid

	pos = pos_a
	visited_mask = visited_mask.at[
		pos[1],pos[0]
	].set(True)
	vstack = jnp.zeros((h*w, 2), dtype=jnp.uint32)
	vstack = vstack.at[:2].set(pos)
	vstack_size = 2

	def _scan_dfs(carry, step):
		(visited_mask, vstack, vstack_size) = carry

		pos = vstack[vstack_size-1]

		neighbors = \
			jnp.minimum(
				jnp.maximum(
					pos + NEIGHBOR_OFFSETS, 0
				), jnp.array([[h,w]])
			).astype(jnp.uint32)

		neighbors_mask = visited_mask.at[
			neighbors[:,1],neighbors[:,0]
		].get()
		n_neighbor_visited = neighbors_mask[:4].sum()
		all_visited = n_neighbor_visited == 4
		all_visited_post = n_neighbor_visited >= 3
		neighbors_mask = neighbors_mask.at[-1].set(~all_visited)

		next_neighbor_idx = jnp.argmax(~neighbors_mask)
		next_neighbor = neighbors[next_neighbor_idx]

		visited_mask = visited_mask.at[
			next_neighbor[1],next_neighbor[0]
		].set(True)

		vstack_size -= all_visited_post

		vstack = vstack.at[vstack_size].set(next_neighbor)
		vstack_size += ~all_visited

		pos = next_neighbor

		return (visited_mask, vstack, vstack_size), None

	max_n_steps = 2*h*w
	(visited_mask, vstack, vstack_size), _ = jax.lax.scan(
		_scan_dfs,
		(visited_mask, vstack, vstack_size),
		jnp.arange(max_n_steps),
		length=max_n_steps
	)

	visited_mask = visited_mask ^ grid
	return visited_mask


@jax.jit
def shortest_path_len(grid, pos_a, pos_b,ignore_value=-1):
	grid = ~component_mask_with_pos(grid, pos_a)

	A = grid_to_graph(grid)
	D = apsp(A, n=A.shape[0])

	if len(pos_b.shape) == 2: # batch eval
		return jax.vmap(_shortest_path_len, in_axes=(None, None, 0, None))(
			grid, pos_a, pos_b, D, ignore_value
		)
	else:
		return _shortest_path_len(grid, pos_a, pos_b, D, ignore_value)


@jax.jit
def _shortest_path_len(grid, pos_a, pos_b, D, ignore_value):
	h,w = grid.shape

	a_idx = pos_a[1]*w + pos_a[0]
	b_idx = pos_b[1]*w + pos_b[0]
	d = D[a_idx][b_idx]

	mhttn_d = jnp.sum(jnp.abs(jnp.maximum(pos_a,pos_b)- jnp.minimum(pos_a,pos_b)))

	impossible = jnp.logical_and(
		d == 1,
		mhttn_d > 1
	)

	return d*(1-impossible)
