import jax.numpy as jnp
import matplotlib.pyplot as plt
from matplotlib.colors import SymLogNorm

from differt2d.geometry import Wall
from differt2d.scene import Scene
from differt2d.utils import received_power

fig, ax = plt.subplots()
scene = Scene.square_scene()
wall = Wall(xys=jnp.array([[0.8, 0.2], [0.8, 0.8]]))
scene = scene.add_objects(wall)
scene.plot(ax, receivers=True)

X, Y = scene.grid(300)
dZ = scene.accumulate_on_receivers_grid_over_paths(
    X,
    Y,
    fun=received_power,
    reduce_all=True,
    grad=True,
)
ndZ = jnp.linalg.norm(dZ, axis=-1)  # Norm of gradient
ndZ = jnp.nan_to_num(ndZ)  # NaN can arise from grad
im = ax.pcolormesh(
    X,
    Y,
    ndZ,
    norm=SymLogNorm(0.5, vmin=ndZ.min(), vmax=ndZ.max()),
    zorder=-1,
)
fig.colorbar(im, ax=ax) # Useful because non-linear scale
plt.show()