import numpy as np
import h5py
import os
from devito.logger import info
from devito import TimeFunction, clear_cache
from examples.seismic.acoustic import AcousticWaveSolver
from examples.seismic import Model, RickerSource, Receiver, TimeAxis
from math import floor
from scipy.interpolate import griddata
import argparse

parser = argparse.ArgumentParser(description='')
parser.add_argument('--data_path', dest='data_path', type=str, default='/home/ec2-user/data', help='raw data path')
parser.add_argument('--save_dir', dest='save_dir', type=str, default='/home/ec2-user/data', help='saving directory')
args = parser.parse_args()
data_path  = args.data_path
save_dir = args.save_dir

origin = (0., 0.)
spacing=(7.5, 7.5)
tn=1100.
nbpml=40
# Define your vp in km/sec (x, z)

vp = np.fromfile(os.path.join(data_path, 'vp_marmousi_bi'),
            dtype='float32', sep="")
vp = np.reshape(vp, (1601, 401))
# vp = vp[400:1401, 0:401]
shape=[401, 301]

values = np.zeros([vp.shape[0]*vp.shape[1], ])
points = np.zeros([vp.shape[0]*vp.shape[1], 2])

k = 0
for indx in range(0, vp.shape[0]):
    for indy in range(0, vp.shape[1]):
        values[k] = vp[indx, indy]
        points[k, 0] = indx
        points[k, 1] = indy

        k = k + 1


# nx, ny = shape[0], shape[1]
X, Y = np.meshgrid(np.array(np.linspace(1000, 1287, shape[0])), np.array(np.linspace(120, 232, shape[1])))

int_vp = griddata(points, values, (X, Y), method='cubic')
int_vp = np.transpose(int_vp)
vp = int_vp

# create model
model = Model(origin, spacing, shape, 2, vp, nbpml=nbpml)
# Derive timestepping from model spacing
dt = model.critical_dt
t0 = 0.0
nt = int(1 + (tn-t0) / dt)  # Number of timesteps
time = np.linspace(t0, tn, nt)  # Discretized time axis


datasize0 = int(np.shape(range(0, shape[0], 4))[0])
datasize1 = int(np.shape(range(100, nt, 20))[0])
datasize = datasize0*datasize1


strTrainA = os.path.join(save_dir, 'Wavefield_Marmousi_pml_401x301_1000-1287_120-232_4k_20kp100_A_train.hdf5')
strTrainB = os.path.join(save_dir, 'Wavefield_Marmousi_pml_401x301_1000-1287_120-232_4k_20kp100_B_train.hdf5')
dataset_train = "train_dataset"

file_trainA = h5py.File(strTrainA, 'w-')
datasetA = file_trainA.create_dataset(dataset_train, (datasize, shape[0]+2*nbpml, shape[1]+2*nbpml))

file_trainB = h5py.File(strTrainB, 'w-')
datasetB = file_trainB.create_dataset(dataset_train, (datasize, shape[0]+2*nbpml, shape[1]+2*nbpml))

num_rec = 601
rec_samp = np.linspace(0., model.domain_size[0], num=num_rec);
rec_samp = rec_samp[1]-rec_samp[0]


time_range = TimeAxis(start=t0, stop=tn, step=dt)
src = RickerSource(name='src', grid=model.grid, f0=0.025, time_range=time_range, space_order=1, npoint=1)
src.coordinates.data[0, :] = np.array([1*spacing[0], 2*spacing[1]]).astype(np.float32)

rec = Receiver(name='rec', grid=model.grid, time_range=time_range, npoint=num_rec)
rec.coordinates.data[:, 0] = np.linspace(0., model.domain_size[0], num=num_rec)
rec.coordinates.data[:, 1:] = src.coordinates.data[0, 1:]

solverbad = AcousticWaveSolver(model, source=src, receiver=rec, kernel='OT2', isic=True,
        space_order=2, freesurface=False)
solvergood = AcousticWaveSolver(model, source=src, receiver=rec, kernel='OT2', isic=True,
        space_order=20, freesurface=False)

ulocgood = TimeFunction(name="u", grid=model.grid, time_order=2, space_order=20, save=nt)
ulocbad = TimeFunction(name="u", grid=model.grid, time_order=2, space_order=2, save=nt)


kk = 0
for xsrc in range(0, shape[0], 4):

    clear_cache()

    ulocgood.data.fill(0.)
    ulocbad.data.fill(0.)

    src.coordinates.data[0, :] = np.array([xsrc*spacing[0], 2*spacing[1]]).astype(np.float32)
    rec.coordinates.data[:, 0] = np.linspace(0., model.domain_size[0], num=num_rec)
    rec.coordinates.data[:, 1:] = src.coordinates.data[0, 1:]

    _, ulocgood, _ = solvergood.forward(m=model.m, src=src, time=nt-1, save=True)
    _, ulocbad, _  = solverbad.forward(m=model.m, src=src, time=nt-1, save=True)

    datasetA[kk:(kk+datasize1), :, :] = np.array(ulocgood.data[range(100, nt, 20), :, :])
    datasetB[kk:(kk+datasize1), :, :] = np.array(ulocbad.data[range(100, nt, 20), :, :])

    kk = kk + datasize1



file_trainA.close()
file_trainB.close()

[((12, 9, 12, 48), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((27, 5, 27, 32), 'numpy.reshape', 'np.reshape', ({(27, 16, 27, 18): 'vp', (27, 20, 27, 31): '(1601, 401)'}, {}), '(vp, (1601, 401))', True, 'import numpy as np\n'), ((31, 9, 31, 46), 'numpy.zeros', 'np.zeros', ({(31, 18, 31, 45): '[vp.shape[0] * vp.shape[1]]'}, {}), '([vp.shape[0] * vp.shape[1]])', True, 'import numpy as np\n'), ((32, 9, 32, 47), 'numpy.zeros', 'np.zeros', ({(32, 18, 32, 46): '[vp.shape[0] * vp.shape[1], 2]'}, {}), '([vp.shape[0] * vp.shape[1], 2])', True, 'import numpy as np\n'), ((47, 9, 47, 57), 'scipy.interpolate.griddata', 'griddata', (), '', False, 'from scipy.interpolate import griddata\n'), ((48, 9, 48, 29), 'numpy.transpose', 'np.transpose', ({(48, 22, 48, 28): 'int_vp'}, {}), '(int_vp)', True, 'import numpy as np\n'), ((52, 8, 52, 57), 'examples.seismic.Model', 'Model', (), '', False, 'from examples.seismic import Model, RickerSource, Receiver, TimeAxis\n'), ((57, 7, 57, 30), 'numpy.linspace', 'np.linspace', ({(57, 19, 57, 21): 't0', (57, 23, 57, 25): 'tn', (57, 27, 57, 29): 'nt'}, {}), '(t0, tn, nt)', True, 'import numpy as np\n'), ((65, 12, 65, 110), 'os.path.join', 'os.path.join', ({(65, 25, 65, 33): 'save_dir', (65, 35, 65, 109): '"""Wavefield_Marmousi_pml_401x301_1000-1287_120-232_4k_20kp100_A_train.hdf5"""'}, {}), "(save_dir,\n    'Wavefield_Marmousi_pml_401x301_1000-1287_120-232_4k_20kp100_A_train.hdf5')", False, 'import os\n'), ((66, 12, 66, 110), 'os.path.join', 'os.path.join', ({(66, 25, 66, 33): 'save_dir', (66, 35, 66, 109): '"""Wavefield_Marmousi_pml_401x301_1000-1287_120-232_4k_20kp100_B_train.hdf5"""'}, {}), "(save_dir,\n    'Wavefield_Marmousi_pml_401x301_1000-1287_120-232_4k_20kp100_B_train.hdf5')", False, 'import os\n'), ((69, 14, 69, 40), 'h5py.File', 'h5py.File', ({(69, 24, 69, 33): 'strTrainA', (69, 35, 69, 39): '"""w-"""'}, {}), "(strTrainA, 'w-')", False, 'import h5py\n'), ((72, 14, 72, 40), 'h5py.File', 'h5py.File', ({(72, 24, 72, 33): 'strTrainB', (72, 35, 72, 39): '"""w-"""'}, {}), "(strTrainB, 'w-')", False, 'import h5py\n'), ((76, 11, 76, 61), 'numpy.linspace', 'np.linspace', (), '', True, 'import numpy as np\n'), ((80, 13, 80, 49), 'examples.seismic.TimeAxis', 'TimeAxis', (), '', False, 'from examples.seismic import Model, RickerSource, Receiver, TimeAxis\n'), ((81, 6, 81, 105), 'examples.seismic.RickerSource', 'RickerSource', (), '', False, 'from examples.seismic import Model, RickerSource, Receiver, TimeAxis\n'), ((84, 6, 84, 82), 'examples.seismic.Receiver', 'Receiver', (), '', False, 'from examples.seismic import Model, RickerSource, Receiver, TimeAxis\n'), ((85, 29, 85, 79), 'numpy.linspace', 'np.linspace', (), '', True, 'import numpy as np\n'), ((88, 12, 89, 41), 'examples.seismic.acoustic.AcousticWaveSolver', 'AcousticWaveSolver', (), '', False, 'from examples.seismic.acoustic import AcousticWaveSolver\n'), ((90, 13, 91, 42), 'examples.seismic.acoustic.AcousticWaveSolver', 'AcousticWaveSolver', (), '', False, 'from examples.seismic.acoustic import AcousticWaveSolver\n'), ((93, 11, 93, 89), 'devito.TimeFunction', 'TimeFunction', (), '', False, 'from devito import TimeFunction, clear_cache\n'), ((94, 10, 94, 87), 'devito.TimeFunction', 'TimeFunction', (), '', False, 'from devito import TimeFunction, clear_cache\n'), ((25, 17, 25, 58), 'os.path.join', 'os.path.join', ({(25, 30, 25, 39): 'data_path', (25, 41, 25, 57): '"""vp_marmousi_bi"""'}, {}), "(data_path, 'vp_marmousi_bi')", False, 'import os\n'), ((100, 4, 100, 17), 'devito.clear_cache', 'clear_cache', ({}, {}), '()', False, 'from devito import TimeFunction, clear_cache\n'), ((106, 33, 106, 83), 'numpy.linspace', 'np.linspace', (), '', True, 'import numpy as np\n'), ((45, 28, 45, 61), 'numpy.linspace', 'np.linspace', ({(45, 40, 45, 44): '1000', (45, 46, 45, 50): '1287', (45, 52, 45, 60): 'shape[0]'}, {}), '(1000, 1287, shape[0])', True, 'import numpy as np\n'), ((45, 73, 45, 104), 'numpy.linspace', 'np.linspace', ({(45, 85, 45, 88): '120', (45, 90, 45, 93): '232', (45, 95, 45, 103): 'shape[1]'}, {}), '(120, 232, shape[1])', True, 'import numpy as np\n'), ((82, 29, 82, 67), 'numpy.array', 'np.array', ({(82, 38, 82, 66): '[1 * spacing[0], 2 * spacing[1]]'}, {}), '([1 * spacing[0], 2 * spacing[1]])', True, 'import numpy as np\n'), ((105, 33, 105, 74), 'numpy.array', 'np.array', ({(105, 42, 105, 73): '[xsrc * spacing[0], 2 * spacing[1]]'}, {}), '([xsrc * spacing[0], 2 * spacing[1]])', True, 'import numpy as np\n')]