import toolbox
import numpy as np
import pylab

#initialise
data, params = toolbox.initialise('prepro.su')
#AGC
toolbox.agc(data, None, **params)


#select the velocities we want to check
params['velocities'] = np.arange(1000, 6000, 100)
params['smoother'] = 25
params['smute'] = 150

#calculate our analysis locations
cdps = np.unique(data['cdp'])
sg_centers = np.arange(800, 3000, 200)

#loop over the supergathers
for sg in sg_centers:
	gathers = range(sg-5, sg+6)
	mask = toolbox.build_mask(data['cdp'], gathers)
	supergather = data[mask]
	supergather.sort(order=['offset'])

	#display it
	toolbox.semb(supergather, **params)

#this is out holding dictionary - we'll put our velocity solutions in it.
vels = {}
vels[753]= (2456.0, 0.153), (2772.1, 0.413), (3003.2, 0.612), (3076.1, 0.704), (3270.7, 1.056), (3367.9, 1.668), (3538.2, 2.204), (3671.9, 3.566), (3915.1, 5.908), 
vels[800]= (2456.0, 0.153), (2772.1, 0.413), (3003.2, 0.612), (3076.1, 0.704), (3270.7, 1.056), (3367.9, 1.668), (3538.2, 2.204), (3671.9, 3.566), (3915.1, 5.908), 
vels[1000]= (2443.9, 0.122), (2662.7, 0.413), (2942.4, 0.980), (3027.5, 1.668), (3890.8, 5.893), 
vels[1200]= (2370.9, 0.138), (3051.8, 0.765), (3367.9, 1.209), (3720.5, 1.730), (4012.3, 2.602), (4292.0, 5.862), (4158.3, 3.689), (2711.4, 0.459), 
vels[1400]= (2419.5, 0.092), (2820.8, 0.520), (3112.6, 1.026), (3392.2, 1.668), (3562.5, 2.128), (3720.5, 2.755), (3830.0, 3.367), (4085.3, 5.939), 
vels[1600]= (2285.8, 0.107), (2662.7, 0.520), (3173.4, 1.087), (3501.7, 1.730), (3647.6, 2.526), (3781.3, 3.184), (4097.5, 5.862), 
vels[1800]= (2322.3, 0.153), (2638.4, 0.566), (3015.3, 1.133), (3355.8, 1.806), (3538.2, 2.556), (3708.4, 3.260), (4000.2, 5.893), 
vels[2000]= (2346.6, 0.199), (2845.1, 0.842), (3039.6, 1.347), (3185.5, 1.638), (3392.2, 3.276), (3963.7, 5.923), (3270.7, 2.510), 
vels[2200]= (2066.9, 0.138), (2480.3, 0.367), (2832.9, 0.765), (3100.4, 1.408), (3258.5, 1.760), (3513.8, 3.061), (3975.9, 5.939), (3416.6, 2.571), 
vels[2400]= (2298.0, 0.168), (2553.3, 0.582), (2735.7, 0.781), (3064.0, 1.301), (3149.1, 2.327), (3392.2, 3.092), (3975.9, 5.985), 
vels[2600]= (2176.4, 0.184), (2529.0, 0.582), (3003.2, 1.347), (3282.8, 1.806), (3465.2, 2.480), (3574.6, 3.184), (4036.7, 5.847), (3902.9, 4.577), 
vels[2800]= (2237.2, 0.168), (2747.8, 0.811), (3197.7, 1.286), (3428.7, 2.128), (3842.1, 3.857), (4012.3, 5.908), (2468.2, 0.490),
vels[3056]= (2237.2, 0.168), (2747.8, 0.811), (3197.7, 1.286), (3428.7, 2.128), (3842.1, 3.857), (4012.3, 5.908), (2468.2, 0.490),
	
params['cdp'] = cdps
params['vels'] = toolbox.build_vels(vels, **params)
np.array(params['vels']).tofile('vels_full.bin')

pylab.imshow(params['vels'].T , aspect='auto')
pylab.colorbar()

pylab.show()