"""
MPI forward modeling on the SEG/Chevron 2014 2D benchmark velocity model.

With MPI enabled, each rank owns a slab of the grid and only needs the part of
the velocity model covering that slab. `WaveModel` accepts, in place of an
array, any model volume that can be sliced: it then reads from it, on every
rank, only the window that rank owns. The full model is therefore never read
from disk nor held in memory, which is what makes large models tractable.

Here that volume is the SEG-Y file itself, read a range of traces at a time.

Run with, e.g.::

    DEVITO_MPI=diag2 mpiexec -n 4 python examples/mpi_modeling.py [path.segy]

The script also runs unchanged on a single process.
"""

import sys

import numpy as np

from pysegy import plot_sdata, segy_scan

from devito import configuration
from devito.logger import info
from devito.mpi import MPI

from examples.seismic import Receiver

from recipes import WaveModel
from recipes.isotropic_acoustic import AcousticIsotropic


# Smooth starting Vp model of the SEG/Chevron 2014 elastic benchmark
VP_FILE = "/Users/mathiaslouboutin/data/ChevronGOM/SEG14.Vpsmoothstarting.segy"

SPACE_ORDER = 8
NBL = 40
F0 = 0.010
RECORD_TIME = 3000.

# Marine acquisition of the benchmark: a streamer of 321 hydrophones, 25 m
# apart and up to 8 km behind the source, all of them 15 m deep
N_REC = 321
REC_SPACING = 25.
DEPTH = 15.

FIGURE_FILE = "./mpi_modeling.png"


class SegyVelocity:
    """
    A velocity model stored as a SEG-Y file, read one window at a time.

    `segy_scan` indexes the file into a single record, a model holding no
    gathers, without reading any trace data. `velocity[i0:i1, j0:j1]` then
    reads the traces `i0:i1` of that record and keeps the samples `j0:j1`,
    which is all `WaveModel` needs to initialize its parameters: it only ever
    asks for the window owned by the calling rank.

    The grid geometry (shape, spacing and origin, in m) comes from the scan,
    and velocities are converted from m/s to the km/s the recipes use.
    """

    def __init__(self, path, scale=1e-3):
        scan = segy_scan(path, keys=['GroupX'])
        if len(scan) != 1:
            raise ValueError(f"{path} scans as {len(scan)} gathers, this "
                             "expects a model")

        self.record = scan[0]
        self.scale = scale

        header = scan.fileheader.bfh
        x0, x1 = self.record.summary['GroupX']

        self.shape = (self.record.ntraces, header.ns)
        self.spacing = ((x1 - x0) / (self.shape[0] - 1), header.dt / 1000)
        self.origin = (x0, 0.)

    def __getitem__(self, window):
        traces, samples = window
        info(f"Rank {MPI.COMM_WORLD.rank}: reading traces "
             f"{traces.start}:{traces.stop} and samples "
             f"{samples.start}:{samples.stop} of {self.shape}")
        # The record holds its traces as columns, and reads those indexed
        return self.scale * self.record.data[samples, traces].T


def run(path=VP_FILE, record_time=RECORD_TIME):
    """
    Model a shot on the benchmark model, with a rank-local model setup.
    """
    velocity = SegyVelocity(path)
    model = WaveModel(velocity.origin, velocity.spacing, velocity.shape,
                      SPACE_ORDER, velocity, nbl=NBL, bcs="damp")

    nt = int(record_time / model.critical_dt)
    if model.grid.distributor.myrank == 0:
        info(f"{velocity.shape} model, {velocity.spacing} m spacing, on a "
             f"{tuple(int(n) for n in model.grid.shape)} grid over "
             f"{model.grid.distributor.nprocs} rank(s), "
             f"dt={model.critical_dt:.3f} ms, {nt} steps")

    solver = AcousticIsotropic(model, {'nt': nt, 'space_order': SPACE_ORDER,
                                       'f0': F0})

    # Shoot from the middle of the line, into a streamer towed behind
    source_x = model.domain_size[0] / 2
    solver.src.coordinates.data[0, :] = np.array([source_x, DEPTH])
    solver.rec = Receiver(name='rec', grid=model.grid, npoint=N_REC,
                          time_range=solver.src.time_range,
                          interpolation=solver.interpolation, r=solver.rinterp,
                          coordinates=np.stack(
                              [source_x + REC_SPACING * np.arange(N_REC),
                               np.full(N_REC, DEPTH)], axis=1))

    solver.forward(save=False)

    # Receivers are distributed too: gather the shot record and the receiver
    # positions onto rank 0, the only rank that then has anything to plot
    shot = solver.rec.data_gather(rank=0)
    coordinates = solver.rec.coordinates.data_gather(rank=0)
    if shot is None:
        return

    plot_sdata(shot, spacing=(1e-3 * model.critical_dt,
                              np.median(np.diff(coordinates[:, 0]))),
               save=FIGURE_FILE)
    info(f"Wrote {FIGURE_FILE} with the {shot.shape} shot record")


if __name__ == "__main__":
    # MPI has to be enabled before the first Grid is created. `DEVITO_MPI`
    # takes precedence, so the mode (`basic`, `diag2`, ...) can be set from
    # the command line
    if not configuration['mpi']:
        configuration['mpi'] = 'basic'

    run(sys.argv[1] if len(sys.argv) > 1 else VP_FILE)