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

if "blob_fmax_1.5" not in p.simulations.list():
    p.add_to_project(
        sn.UnstructuredMeshSimulationConfiguration(
            name="blob_fmax_1.5", unstructured_mesh=m2m, event_configuration=ec
        )
    )

for name in ["blob_fmax", "blob_fmax_1.5", "homo_fmax_1.5"]:
    p.simulations.launch(
        simulation_configuration=name,
        events=p.events.list(),
        site_name=SALVUS_FLOW_SITE_NAME,
        ranks_per_job=1,
    )
    p.simulations.query(block=True)

p.viz.waveforms(
    event="event_0000",
    receiver_name="XX.XX",
    data=["blob_fmax", "blob_fmax_1.5", "homo_fmax_1.5"],
    receiver_field="phi",
)

# Spherical coordinates
lat = np.linspace(-5.0, 5.0, 101)
lon = np.linspace(-5.0, 5.0, 101)
dep = np.linspace(0.0, 660e3, 101)

# Gaussian
xx, yy, zz = np.meshgrid(lat, lon, dep, indexing="xy")
g = np.zeros_like(xx)
for lats in lat[25:-25:25]:
    for lons in lon[::25]:
        g += np.exp(-(((xx - lats) ** 2 + (yy - lons) ** 2) / (2 * 0.5**2)))

# Xarray dataset
ds = xr.Dataset(
    coords={"latitude": lat, "longitude": lon, "depth": dep},
    data_vars={
        "vp": (["latitude", "longitude", "depth"], g * 20, {"units": "%"}),
    },
    attrs={
        "geospatial_lon_units": "degrees",
        "geospatial_lat_units": "degrees_north",
        "geospatial_vertical_units": "m",
    },
)

ds_mantle = ds.copy()
ds_mantle["vp"] *= -1

# Salvus wrapper.
mc = sn.model.volume.seismology.CrustalModel(name="crust", data=ds)
mm = sn.model.volume.seismology.MantleModel(name="mantle", data=ds)

# Plot depth slice.
mm.ds.vp.isel(depth=0).plot()

d = sn.domain.dim3.SphericalChunkDomain(
    lat_center=0.0,
    lon_center=0.0,
    lat_extent=5.0,
    lon_extent=5.0,
    radius_in_meter=6371e3,
)

# Add mantle and crustal model.
p = sn.Project.from_domain("proj_3d", d, True)
for m in [mc, mm]:
    p.add_to_project(m, overwrite=True)

s = sn.simple_config.source.seismology.SideSetVectorPoint3D(
    fr=1e10,
    ft=0.0,
    fp=0.0,
    latitude=-2.5,
    longitude=-2.5,
    depth_in_m=0.0,
    side_set_name="r1",
)

r = sn.simple_config.receiver.seismology.SideSetPoint3D(
    latitude=2.5,
    longitude=2.5,
    fields=["velocity"],
    station_code="XX",
    network_code="XX",
    side_set_name="r1",
)

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

p_min = 50.0

# Common event configuration.
ec = sn.EventConfiguration(
    wavelet=sn.simple_config.stf.Ricker(center_frequency=1 / (2 * p_min)),
    waveform_simulation_configuration=sn.WaveformSimulationConfiguration(
        end_time_in_seconds=600.0
    ),
)

# Heterogeneous - meshed at f_max.
s0 = sn.SimulationConfiguration(
    name="blobs",
    tensor_order=2,
    max_depth_in_meters=660e3,
    min_period_in_seconds=p_min,
    event_configuration=ec,
    model_configuration=sn.ModelConfiguration(
        background_model=sn.model.background.one_dimensional.BuiltIn(
            name="prem_iso_one_crust"
        ),
        volume_models=["mantle", "crust"],
    ),
    elements_per_wavelength=2.0,
)

# 1D only - meshed at a slightly lower resolution.
s1 = sn.SimulationConfiguration(
    name="1d",
    tensor_order=2,
    max_depth_in_meters=660e3,
    min_period_in_seconds=p_min,
    event_configuration=ec,
    model_configuration=sn.ModelConfiguration(
        background_model=sn.model.background.one_dimensional.BuiltIn(
            name="prem_iso_one_crust"
        ),
    ),
    elements_per_wavelength=1.5,
)

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

p.viz.nb.simulation_setup("blobs")

# Interpolate.
m2m = interpolate_mesh_to_mesh(
    p.simulations.get_mesh("blobs"),
    p.simulations.get_mesh("1d"),
    use_layers=True,
    use_1d_vertical_coordinate=True,
)

# Assert mesh is different.
assert m2m.nelem != p.simulations.get_mesh("blobs").nelem

# Visualize.
m2m

p.add_to_project(
    sn.UnstructuredMeshSimulationConfiguration(
        name="blobs_reinterp", unstructured_mesh=m2m, event_configuration=ec
    )
)

for name in ["blobs", "blobs_reinterp", "1d"]:
    p.simulations.launch(
        simulation_configuration=name,
        events=p.events.list(),
        site_name=SALVUS_FLOW_SITE_NAME,
        ranks_per_job=1,
    )
    p.simulations.query(block=True)

p.viz.waveforms(
    event="event_0000",
    receiver_name="XX.XX",
    data=["blobs", "blobs_reinterp", "1d"],
    receiver_field="velocity",
)