# This file is part of j-Wave.
#
# j-Wave is free software: you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation, either
# version 3 of the License, or (at your option) any later version.
#
# j-Wave is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with j-Wave. If not, see <https://www.gnu.org/licenses/>.

from jax import jit

from jwave import FourierSeries
from jwave.acoustics.time_varying import simulate_wave_propagation
from jwave.geometry import Domain, Medium, TimeAxis
from jwave.utils import load_image_to_numpy

# Simulation parameters
N, dx = (256, 256), (0.1e-3, 0.1e-3)
domain = Domain(N, dx)
medium = Medium(domain=domain, sound_speed=1500.0)
time_axis = TimeAxis.from_medium(medium, cfl=0.3, t_end=0.8e-05)

# Initial pressure field
p0 = load_image_to_numpy("docs/assets/images/jwave.png", image_size=N) / 255.0
p0 = FourierSeries(p0, domain)


# Compile and run the simulation
@jit
def solver(medium, p0):
    return simulate_wave_propagation(medium, time_axis, p0=p0)


pressure = solver(medium, p0)

if __name__ == "__main__":
    from matplotlib import pyplot as plt

    from jwave.utils import show_field

    _ = show_field(pressure[250])
    plt.title(f"Pressure at time t={time_axis.to_array()[250]}")
    plt.savefig("docs/assets/images/readme_example_reconimage.png")
