import time
import sys

import numpy as np

from pysit import *
from pysit.gallery import marmousi

from pysit.util.parallel import ParallelWrapShot

from mpi4py import MPI

if __name__ == '__main__':

	comm = MPI.COMM_WORLD
	size = comm.Get_size()
	rank = comm.Get_rank()

	# Wrapper for handling parallelism over shots
	pwrap = ParallelWrapShot(comm=MPI.COMM_WORLD)

	if rank == 0:
		ttt = time.time()

	#	Load or generate true wave speed
	C, C0, m, d = marmousi(patch='mini_square')

	# Set up shots

	Nshots = size
	sys.stdout.write("{0}: {1}\n".format(rank, Nshots / size))

	shots = equispaced_acquisition(m,
	                               RickerWavelet(10.0),
	                               sources=Nshots,
	                               source_depth=20.0,
	                               receivers='max',
	                               parallel_shot_wrap=pwrap,
	                               )

	# Define and configure the wave solver
	trange = (0.0,3.0)

	solver = ConstantDensityAcousticWave(m,
	                                     formulation='scalar',
	                                     model_parameters={'C': C},
	                                     spatial_accuracy_order=6,
	                                     trange=trange,
	                                     use_cpp_acceleration=True)

	# Generate synthetic Seismic data
	sys.stdout.write('Generating data...')
	base_model = solver.ModelParameters(m,{'C': C})
	tt = time.time()
	generate_seismic_data(shots, solver, base_model)
	data_gen_time = time.time()-tt
	sys.stdout.write('{1}:Data generation: {0}s\n'.format(data_gen_time,rank))
	sys.stdout.flush()

	comm.Barrier()

	if rank == 0:
		tttt = time.time()-ttt
		sys.stdout.write('Total wall time: {0}\n'.format(tttt))
		sys.stdout.write('Total wall time/shot: {0}\n'.format(tttt/Nshots))

	objective = TemporalLeastSquares(solver,
		                             parallel_wrap_shot=pwrap)

	# Define the inversion algorithm
	invalg = LBFGS(objective, memory_length=10)

	initial_value = solver.ModelParameters(m, {'C': C0})

	# Execute inversion algorithm
	print('Running LBFGS...')
	tt = time.time()

	nsteps = 10

	status_configuration = {'value_frequency'           : 1,
	                        'objective_frequency'       : 1}

	result = invalg(shots, initial_value, nsteps,
	                line_search='backtrack',
	                status_configuration=status_configuration, 
	                verbose=True)

	print 'Run time:  {0}s'.format(time.time()-tt)


	if rank == 0:

		model = result.C.reshape(m.shape(as_grid=True))

		vals = list()
		for k,v in invalg.objective_history.items():
			vals.append(v)
		obj_vals = np.array(vals)

		from scipy.io import savemat

		out = {'result': model,
			   'true': C.reshape(m.shape(as_grid=True)),
			   'initial': C0.reshape(m.shape(as_grid=True)),
			   'obj': obj_vals
			  }
		savemat('marm_recon.mat',out)
