#! /bin/usr/env python

import argparse

import numpy
import numpy.linalg
import random
from mpi4py import MPI

from VelocityModel import VelocityModel
from AcousticWave2D import AcousticWave2D
from ShotManager import ShotManager

from scipy.ndimage.filters import gaussian_filter
from TaskFarm import Master, Worker
import scipy.optimize

global comm
comm = MPI.COMM_WORLD
rank = comm.Get_rank()

if rank == 0:
    parser = argparse.ArgumentParser()
    parser.add_argument('--source', '-s',
                        help=('Specify the source file. ' +
                              'Only SEG-Y rev-1 supported.'))
    parser.add_argument('--model', '-m',
                        help=('Starting model. Only RSF ' +
                              'formatted files supported.'))
    parser.add_argument('shots', action='append', nargs='+',
                        help=('List of shots/receiver data. ' +
                              'Only SEG-Y rev-1 supported.'))

    cmd_args = parser.parse_args()
if rank != 0:
    worker = Worker(comm, AcousticWave2D)
    worker.run()
else:
    global model
    model = VelocityModel(cmd_args.model, subsample=15)

    # As this is just a synthetic example we take the input
    # as the exact model, and apply a gaussian smooth to
    # create the 'initial' dataset.
    vp_exact = numpy.copy(model.vp)
    model.vp = gaussian_filter(vp_exact, sigma=7)

    global shots
    shots = ShotManager(comm, cmd_args.shots[0])

    global source
    if cmd_args.source:
        source = ShotManager(comm, [cmd_args.source])
    else:
        source = None

    def function_fwi(x):
        model.vp = numpy.reshape(numpy.sqrt(1/x), model.vp.shape)
        master = Master(comm, model, shots, source)

        nshots = shots.get_number_of_shots()
        nworkers = comm.Get_size() - 1

        worklist = random.sample(range(nshots), nworkers)

        f, g = master.run(worklist)

        return f, .01*g.reshape(-1)/(numpy.abs(numpy.amax(g)))

    scipy.optimize.minimize(function_fwi, (model.vp**(-2)).reshape(-1),
                            method='CG', jac=True,
                            options={"maxiter": 1, "disp": True})

    true_residule = numpy.linalg.norm(model.vp - vp_exact)

    model.write(cmd_args.model.replace('.rsf', '-opt.rsf'))
    model.export_vtk(cmd_args.model.replace('.rsf', '-opt.vtk'))
