import numpy as np
from smii.modeling.wavelets.wavelets import ricker
from smii.modeling.propagators.propagators import Scalar2D

def setup_seam(seam_model_path):
    """Prepare the SEAM models and dataset metadata.

    Args:
        seam_model_path: The path to the file containing the true SEAM model

    Returns:
        A dictionary containing SEAM-related arrays and properties
    """
    nx = 64
    nz = 64
    model_true = (np.fromfile(seam_model_path, dtype=np.float32)
                  .reshape([nx, nz]).T)

    dx = 100
    dt = 0.005
    num_sources = 13
    dsrc = 5
    num_receivers = 64
    drec = 1
    nt = int(12 / dt) # 12 seconds
    sources = []
    for i in range(num_sources):
        source = {}
        source['amplitude'] = ricker(1, nt, dt, 1.0)[np.newaxis, :]
        source['locations'] = np.array([0, dsrc/2 + i*dsrc], np.int)
        sources.append(source)
    receivers_x = np.zeros([num_receivers, 2], np.int)
    receivers_x[:, 0] = 0
    receivers_x[:, 1] = np.arange(0, num_receivers*drec, drec)

    return {'model_true': model_true,
            'dx': dx,
            'dt': dt,
            'sources': sources,
            'receivers_x': receivers_x}