How Much Information Does a Shot Gather Contain? / 一个共炮点道集蕴含了多少信息

"""
Deepwave Easy API (Refactored)
==============================
基于物理坐标 (米) 和纯面向对象架构的 2D 地震波正演封装。
核心设计理念：数据（物理模型、观测系统）与计算（仿真器）彻底分离。
"""

import numpy as np
import torch
import deepwave
from deepwave import elastic, scalar
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# =============================================================================
# 0. 信号工具 (Signal Utilities)
# =============================================================================


def ricker(f0, nt, dt, delay=None):
    """
    生成 Ricker 子波。

    参数
    ----------
    f0 : float, 主频 (Hz)
    nt : int, 采样点数
    dt : float, 采样间隔 (s)
    delay : float, 延迟时间 (s)，默认 1.5/f0
    """
    if delay is None:
        delay = 1.5 / f0
    t = np.arange(nt) * dt - delay
    a = (np.pi * f0 * t) ** 2
    return ((1 - 2 * a) * np.exp(-a)).astype(np.float32)


# =============================================================================
# 1. 数据对象层 (Data Objects) - 只存储信息，不处理计算逻辑
# =============================================================================

class PhysicalModel:
    """
    物理地质模型。仅持有介质参数和网格间距，不关心具体的观测系统和网格索引。
    """

    def __init__(self, vp, vs=None, rho=None, dh=10.0, device=None):
        self.device = device or torch.device(
            'cuda' if torch.cuda.is_available() else 'cpu')
        self.dh = float(dh)

        # 统一转为张量并加载到指定设备
        self.vp = self._to_tensor(vp)
        self.vs = self._to_tensor(
            vs) if vs is not None else torch.zeros_like(self.vp)
        self.rho = self._to_tensor(
            rho) if rho is not None else torch.ones_like(self.vp) * 2000.0

        self.nz, self.nx = self.vp.shape
        # [left, right, bottom, top] 用于 matplotlib extent (注意 Z 轴向下)
        self.extent = [0, (self.nx - 1) * self.dh, (self.nz - 1) * self.dh, 0]

    def _to_tensor(self, data):
        if isinstance(data, torch.Tensor):
            return data.to(self.device, dtype=torch.float32)
        return torch.tensor(data, device=self.device, dtype=torch.float32)

    @property
    def lame_parameters(self):
        """返回 Deepwave 弹性波仿真所需的 Lamé 参数"""
        return deepwave.common.vpvsrho_to_lambmubuoyancy(self.vp, self.vs, self.rho)

    def to_cpu(self):
        """将 Vp 提取到 CPU 方便绘图"""
        return self.vp.detach().cpu().numpy()


class Acquisition:
    """
    观测系统配置。仅持有物理坐标（米）、时间步长和子波。
    """

    def __init__(self, src_x, src_z, rec_x, rec_z, wavelet, dt):
        # 将标量深度自动广播匹配所有的 x 坐标
        self.src_x = np.atleast_1d(src_x)
        self.src_z = np.broadcast_to(np.atleast_1d(src_z), self.src_x.shape)

        self.rec_x = np.atleast_1d(rec_x)
        self.rec_z = np.broadcast_to(np.atleast_1d(rec_z), self.rec_x.shape)

        self.wavelet = wavelet  # 形状 (nt,)
        self.dt = float(dt)
        self.nt = len(wavelet)
        self.n_shots = len(self.src_x)
        self.n_recs = len(self.rec_x)


# =============================================================================
# 2. 仿真引擎层 (Simulator) - 处理所有转换、PML 边界和计算逻辑
# =============================================================================

class SeismicSimulator:
    """
    仿真大脑：负责将 PhysicalModel 和 Acquisition 结合，计算网格索引，调用 Deepwave 核心引擎。
    """

    def __init__(self, model: PhysicalModel, acquisition: Acquisition):
        self.model = model
        self.acq = acquisition
        self.device = model.device

    def _get_grid_locations(self, x_phys, z_phys):
        """将物理坐标(m)精确映射为 Deepwave 所需的网格索引张量"""
        iz = np.round(z_phys / self.model.dh).astype(np.int64)
        ix = np.round(x_phys / self.model.dh).astype(np.int64)

        # 边界安全裁剪，防止坐标越界
        iz = np.clip(iz, 0, self.model.nz - 1)
        ix = np.clip(ix, 0, self.model.nx - 1)

        # 构造 Deepwave 需要的形状 (n_shots, n_locations, 2)
        # 这里默认每个炮都使用相同的检波器排列（固定排列）
        locs = torch.zeros(self.acq.n_shots, len(
            ix), 2, dtype=torch.long, device=self.device)
        locs[..., 0] = torch.from_numpy(iz)
        locs[..., 1] = torch.from_numpy(ix)
        return locs

    def run_acoustic(self, pml_width=30, free_surface=True, accuracy=4, **kwargs):
        """
        运行声波仿真 (只使用 Vp)。
        **kwargs 用于接收 forward_callback 等高级参数。
        """
        src_locs = self._get_grid_locations(self.acq.src_x, self.acq.src_z)
        rec_locs = self._get_grid_locations(self.acq.rec_x, self.acq.rec_z)

        # 扩展子波维度以匹配 (n_shots, n_sources_per_shot, nt)
        wav = torch.tensor(self.acq.wavelet, device=self.device).reshape(
            1, 1, -1).expand(self.acq.n_shots, 1, -1)

        # 自由地表处理
        pml = [0, pml_width, pml_width,
               pml_width] if free_surface else [pml_width] * 4

        out = scalar(
            self.model.vp, self.model.dh, self.acq.dt,
            source_amplitudes=wav,
            source_locations=src_locs,
            receiver_locations=rec_locs,
            pml_width=pml,
            accuracy=accuracy,
            **kwargs
        )
        return out[-1]  # 返回 receiver_amplitudes，形状为 (n_shots, n_receivers, nt)

    def run_elastic(self, pml_width=30, free_surface=True, src_comp='y', rec_comp='y', accuracy=4, **kwargs):
        """
        运行弹性波仿真。
        """
        lamb, mu, buoy = self.model.lame_parameters
        src_locs = self._get_grid_locations(self.acq.src_x, self.acq.src_z)
        rec_locs = self._get_grid_locations(self.acq.rec_x, self.acq.rec_z)
        wav = torch.tensor(self.acq.wavelet, device=self.device).reshape(
            1, 1, -1).expand(self.acq.n_shots, 1, -1)

        pml = [0, pml_width, pml_width,
               pml_width] if free_surface else [pml_width] * 4

        src_args = {f'source_amplitudes_{src_comp}': wav,
                    f'source_locations_{src_comp}': src_locs}
        rec_args = {f'receiver_locations_{rec_comp}': rec_locs}

        out = elastic(
            lamb, mu, buoy, self.model.dh, self.acq.dt,
            **src_args, **rec_args,
            pml_width=pml, accuracy=accuracy,
            **kwargs
        )

        idx = -2 if rec_comp in ['y', 'z'] else -1
        return out[idx]

    def record_wavefield(self, mode='acoustic', pml_width=30, free_surface=True, freq_interval=5):
        """
        录制波场传播过程。

        参数
        ----------
        mode : str, 'acoustic' 或 'elastic'
        freq_interval : int, 采样间隔（帧数）。减小该值能提升流畅度，但会增加内存占用。

        返回
        -------
        torch.Tensor, 录制的波场快照，形状为 (n_frames, n_shots, nz, nx)
        """
        wavefield_history = []

        def callback(state):
            # 获取当前波场，为了节省显存，直接 clone 并放到 CPU 上
            wf = state.get_wavefield("wavefield_0")[0].cpu().clone()
            wavefield_history.append(wf)

        if mode == 'acoustic':
            self.run_acoustic(pml_width=pml_width, free_surface=free_surface,
                              forward_callback=callback,
                              callback_frequency=freq_interval)
        elif mode == 'elastic':
            self.run_elastic(pml_width=pml_width, free_surface=free_surface,
                             forward_callback=callback,
                             callback_frequency=freq_interval)
        else:
            raise ValueError("mode must be 'acoustic' or 'elastic'")

        return torch.stack(wavefield_history)


# =============================================================================
# 3. 绘图与可视化工具 (Visualization Utilities)
# =============================================================================

def plot_simulation(simulator: SeismicSimulator, data=None):
    """
    一键可视化模型、观测系统和生成的共炮点道集（Shot Gather）。
    """
    fig, axes = plt.subplots(1, 2, figsize=(15, 6))
    ax1, ax2 = axes

    # --- 1. 速度模型与几何 ---
    m = simulator.model
    a = simulator.acq
    im = ax1.imshow(m.to_cpu(), extent=m.extent, cmap='terrain', aspect='auto')

    # 标出震源和检波器
    ax1.scatter(a.src_x, a.src_z, c='red', marker='*',
                s=150, label='Sources', edgecolors='black')
    ax1.scatter(a.rec_x, a.rec_z, c='blue', marker='v',
                s=20, alpha=0.5, label='Receivers')

    ax1.set_title("Velocity Model (Vp) & Survey Geometry")
    ax1.set_xlabel("Lateral Position (m)")
    ax1.set_ylabel("Depth (m)")
    ax1.legend(loc='lower left')
    plt.colorbar(im, ax=ax1, label='Velocity (m/s)')

    # --- 2. 炮集 (可选) ---
    if data is not None:
        # 取第一炮的数据，转移到 numpy
        d = data[0].detach().cpu().numpy() if isinstance(
            data, torch.Tensor) else data[0]
        # 使用 98% 分位数作为显示裁剪阈值（Clip）
        clip = np.percentile(np.abs(d), 98)

        ax2.imshow(d.T, aspect='auto', cmap='gray', vmin=-clip, vmax=clip,
                   extent=[a.rec_x[0], a.rec_x[-1], a.nt * a.dt, 0])
        ax2.set_title("Synthetic Shot Gather (Shot 0)")
        ax2.set_xlabel("Receiver Position (m)")
        ax2.set_ylabel("Time (s)")

    plt.tight_layout()
    plt.show()


def save_wavefield_animation(wavefields, model: PhysicalModel, filename="wavefield.gif", shot_idx=0):
    """
    将录制的波场张量转化为 GIF 动画并保存。

    参数
    ----------
    wavefields : torch.Tensor, 形状 (n_frames, n_shots, nz, nx)
    model : PhysicalModel, 用于获取坐标范围
    shot_idx : int, 选择录制哪一炮的波场
    """
    fig, ax = plt.subplots(figsize=(10, 6))

    # 提取选定炮的波场数据
    wf_data = wavefields[:, shot_idx, :, :]

    # 选取整个动画序列绝对值最大值的 10% 作为颜色条阈值
    vlimit = torch.max(torch.abs(wf_data)) * 0.1

    im = ax.imshow(wf_data[0].numpy(), cmap='RdBu', aspect='auto',
                   extent=model.extent, vmin=-vlimit, vmax=vlimit)

    ax.set_title(f"Wavefield Propagation (Shot {shot_idx})")
    ax.set_xlabel("Lateral Position (m)")
    ax.set_ylabel("Depth (m)")
    plt.colorbar(im, label="Amplitude")

    def update(frame):
        im.set_array(wf_data[frame].numpy())
        return [im]

    print(f"开始渲染动画，共 {len(wf_data)} 帧...")
    ani = animation.FuncAnimation(fig, update, frames=len(wf_data),
                                  interval=50, blit=True)

    # 使用 pillow 保存 gif
    ani.save(filename, writer='pillow')
    print(f"✅ 动画已成功保存至: {filename}")
    plt.close()


# =============================================================================
# 测试运行块 (Example Usage)
# =============================================================================
if __name__ == "__main__":
    # 1. 建立具有两层的地质模型
    vp_grid = np.ones((150, 300)) * 2000.0
    vp_grid[80:, :] = 3000.0  # 下层高速
    model = PhysicalModel(vp=vp_grid, dh=5.0)  # 网格间距 5m

    # 2. 生成 Ricker 子波并配置观测系统
    dt = 0.001
    wav = ricker(f0=25, nt=600, dt=dt)

    acq = Acquisition(
        src_x=[750], src_z=[10],                   # 地表中点激发
        rec_x=np.linspace(0, 1500, 150), rec_z=10,  # 全排列接收
        wavelet=wav, dt=dt
    )

    # 3. 初始化仿真器
    sim = SeismicSimulator(model, acq)

    # [选项 A] 生成常规地震数据，用于 AI 训练
    print("正在计算共炮点道集...")
    shot_gather = sim.run_acoustic(pml_width=30, free_surface=True)
    plot_simulation(sim, shot_gather)

    # [选项 B] 录制并生成波场动画
    print("正在录制地下波场...")
    wf_history = sim.record_wavefield(mode='acoustic', freq_interval=6)
    save_wavefield_animation(wf_history, model, filename="wavefield_demo.gif")

import os
import numpy as np
from scipy.interpolate import RegularGridInterpolator
import deepwave_easy_api as dwe
import deepwave
import segyio
import matplotlib.pyplot as plt
from tqdm import tqdm
import torch

DH = 1.249  # m


def main():

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Using device: {device}")

    datadir = "../data/"

    f0 = 15
    time = 4000  # ms
    dt = 0.0002

    nt = round(time / dt)

    dsrc_x = 25
    drec_x = 12.5

    data = np.load(os.path.join(datadir, "velmodel.npz"))
    vp = data["vp"]
    vs = data["vs"]
    dens = data["dens"]
    dh = data["dh"]

    print("max_vp", vp.max())
    print("min_vp", vp.min())

    nz, nx = vp.shape

    src_x = np.arange(0, dh * (nx - 1), dsrc_x)
    src_z = np.ones_like(src_x) * 10

    n_shots = len(src_x)
    print("dsrc_x:", dsrc_x)
    print(f"number of sources: {n_shots}")

    rec_x = np.arange(0, dh * (nx - 1), drec_x)
    rec_z = np.ones_like(rec_x) * 10
    n_receivers = len(rec_x)
    print("drec_x:", drec_x)
    print(f"number of receivers: {n_receivers}")

    # output src positions
    data_dir = os.path.abspath(os.path.join(
        os.path.dirname(__file__), "..", "data"))
    os.makedirs(data_dir, exist_ok=True)
    geom_path = os.path.join(data_dir, "geometry.npz")
    np.savez(geom_path, src_x=src_x, src_z=src_z,
             rec_x=rec_x, rec_z=rec_z, dt=dt, nt=nt, f0=f0, dh=dh)
    print(f"Saved geometry to {geom_path}")

    model = dwe.VelocityModel(vp, vs, dens, dh=dh)

    wavelet = dwe.ricker(f0, nt, dt)

    shots_dir = os.path.join(data_dir, "shots")
    os.makedirs(shots_dir, exist_ok=True)

    for i in tqdm(range(n_shots), desc="Generating shots", ncols=100):
        geom = dwe.Geometry(model, src_x[i], src_z[i], rec_x, rec_z, wavelet)
        shot = dwe.elastic2d(model, geom, dt=dt,
                             pml_width=100, pml_freq=f0, accuracy=4)
        np.savez(os.path.join(shots_dir, f"shot_{i:04d}.npz"), shot=shot)
        print(f"Saved shot {i} to {shots_dir}", end="\r", flush=True,)


if __name__ == '__main__':
    main()