# Std import block
import time
import copy

import numpy as np
import matplotlib.pyplot as plt

from pysit import *
from pysit.gallery import marmousi
from pysit.gallery import marmousi2

if __name__ == '__main__':
	#	Load or generate true wave speed
	C, C0, m, d = marmousi(patch='mini_square')
#	C, C0, m, d = marmousi2(patch='mini_square')

	# Set up shots
	shots = equispaced_acquisition(m,
	                               RickerWavelet(10.0),
	                               sources=1,
	                               source_depth=500.0,
	                               source_kwargs={},
	                               receivers='max',
	                               receiver_kwargs={},
	                               )
	
	# 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
	print('Generating data...')
	base_model = solver.ModelParameters(m,{'C': C})
	tt = time.time()
	wavefields = []
	generate_seismic_data(shots, solver, base_model, wavefields=wavefields)
	print 'Data generation: {0}s'.format(time.time()-tt)

	objective = TemporalLeastSquares(solver)
	
	# Define the inversion algorithm
	invalg = LBFGS(objective)
	initial_value = solver.ModelParameters(m,{'C': C0})

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

	nsteps = 5
	status_configuration = {'value_frequency'           : 1, 
	                 'residual_length_frequency' : 1, 
	                 'objective_frequency'       : 1, 
	                 'step_frequency'            : 1,
	                 'step_length_frequency'     : 1,
	                 'gradient_frequency'        : 1, 
	                 'gradient_length_frequency' : 1, 
	                 'run_time_frequency'        : 1, 
	                 'alpha_frequency'           : 1,
	                 }
	
	line_search = 'backtrack'
	
	result = invalg(shots, initial_value, nsteps, 
	                line_search=line_search, 
	                status_configuration=status_configuration, verbose=True)
	                
	print '...run time:  {0}s'.format(time.time()-tt)
	
	obj_vals = np.array([v for k,v in invalg.objective_history.items()])
	
	plt.figure()
	plt.semilogy(obj_vals)
		
	clim = C.min(),C.max()
		
	# Do something to visualize the results
	plt.figure()	
	plt.subplot(3,1,1)
	vis.plot(C0, m, clim=clim)
	plt.title('Initial Model')
	plt.subplot(3,1,2)
	vis.plot(C, m, clim=clim)
	plt.title('True Model')
	plt.subplot(3,1,3)
	vis.plot(result.C, m, clim=clim)
	plt.title('Reconstruction')
	
	plt.figure()	
	plt.subplot(2,1,1)
	vis.plot(result.C-C0, m)
	plt.title('Recon - Initial')
	plt.subplot(2,1,2)
	vis.plot(C-result.C, m)
	plt.title('True-Recon')
	
	plt.show()
	

