import os
import numpy as np
import xarray as xr
import salvus.namespace as sn
from salvus.mesh.tools.transforms import interpolate_mesh_to_mesh

# Run simulations on this site.
SALVUS_FLOW_SITE_NAME = os.environ.get("SITE_NAME", "local")

# Gaussian
x, y = np.linspace(0.0, 1.0, 100), np.linspace(0.0, 1.0, 100)
xx, yy = np.meshgrid(x, y, indexing="xy")
g = np.exp(-(((xx - 0.5) ** 2 + (yy - 0.5) ** 2) / (2 * 0.2**2)))

# Pars
vp = 2 * g + 1
rho = vp / 2

# Xarray dataset
ds = xr.Dataset(
    coords={"x": x, "y": y},
    data_vars={"vp": (["x", "y"], vp), "rho": (["x", "y"], rho)},
)

# Salvus wrapper.
m = sn.model.volume.cartesian.GenericModel(name="blob", data=ds)

# Plot
m.ds.vp.plot()

p = sn.Project.from_volume_model("proj_2d", m, True)

s = sn.simple_config.source.cartesian.ScalarPoint2D(x=0.25, y=0.25, f=1.0)
r = sn.simple_config.receiver.cartesian.Point2D(
    x=0.75, y=0.75, fields=["phi"], station_code="XX", network_code="XX"
)

p.add_to_project(sn.EventCollection.from_sources(sources=s, receivers=r))

f_max = 10.0

# Common event configuration.
ec = sn.EventConfiguration(
    wavelet=sn.simple_config.stf.Ricker(center_frequency=f_max / 2),
    waveform_simulation_configuration=sn.WaveformSimulationConfiguration(
        end_time_in_seconds=1.0
    ),
)

# Heterogeneous - meshed at f_max.
s0 = sn.SimulationConfiguration(
    name="blob_fmax",
    max_frequency_in_hertz=f_max,
    event_configuration=ec,
    model_configuration=sn.ModelConfiguration(
        background_model=None, volume_models="blob"
    ),
    elements_per_wavelength=2.0,
)

# Homogeneous - meshed at 1.5 * f_max.
s1 = sn.SimulationConfiguration(
    name="homo_fmax_1.5",
    max_frequency_in_hertz=1.5 * f_max,
    event_configuration=ec,
    model_configuration=sn.ModelConfiguration(
        background_model=sn.model.background.homogeneous.IsotropicAcoustic(
            rho=rho.min(), vp=vp.min()
        ),
    ),
    elements_per_wavelength=2.0,
)

for s in [s0, s1]:
    p.add_to_project(s, overwrite=True)

# Interpolate.
m2m = interpolate_mesh_to_mesh(
    p.simulations.get_mesh("blob_fmax"),
    p.simulations.get_mesh("homo_fmax_1.5"),
    use_layers=False,
    use_1d_vertical_coordinate=False,
)

# Visualize.
m2m