# Import the numpy package
import numpy as np

# Import pyplot
import matplotlib.pyplot as plt

# Create a function to make the Jacobian
def VSPzo_Jacobian( receiverVector : float, layerVector : float ):
	"""
	Create the model, geometry of sources and receivers, and
	the associated depth-matrix for a zero-offset transmission VSP: 
	source at surface (z=0) and receivers in a vertical borehole

	Args:
		receiverVector (float): A 1D array of floats that holds the depth position of each receiver
		layerVector    (float): A 1D array of floats that holds the depth each layer in the model

	Returns:
		depths (float): A 2D array of floats that holds the geometry of the model
	"""

	# Get the number of layers and receivers
	numberOfReceivers = len( receiverVector )
	numberOfLayers    = len( layerVector    )

	# Check if the number of layers is larger than the number of receivers (model parameters > data points)
	if ( numberOfLayers > numberOfReceivers ):

		# Print info
		print( f"Number of layers ({numberOfLayers}) is larger than the number of receivers({numberOfReceivers})" )

		# Create a matrix of depths
		depths = np.zeros( ( numberOfReceivers, numberOfLayers - 1 ) )

		# Loop over the number of receivers
		for iRec in range( numberOfReceivers ):

			# Loop over the number of layers
			for iLay in range( numberOfLayers ):

				# Check if the depth of the layer is larger than the current receiver depth
				if layerVector[ iLay ] > receiverVector[ iRec ]:

					# Set the depth
					depths[ iRec, iLay ] = receiverVector[ iRec ] - layerVector[ iLay - 1 ]

					# Break off here
					break

				# Check if the depth of the layer is smaller than or equal to the current receiver depth
				else :

					# Check if this is the first layer
					if iLay == 0:

						# Set the depth to the first layer depth
						depths[ iRec, iLay ] = layerVector[ iLay ]

					# Check if this is not the first layer
					else:

						# Set the depth equal to the layer depth difference
						depths[ iRec, iLay ] = layerVector[ iLay ] - layerVector[ iLay - 1 ]

	# Check if the number of layers is smaller than or equal to the number of receivers (model parameters > data points)
	else:

		# Print info
		print( f"Number of layers ({numberOfLayers}) is smaller than or equal to the number of receivers({numberOfReceivers})" )

		# Create a matrix of depths
		depths = np.zeros( ( numberOfReceivers, numberOfLayers ) )

		# Set the layer index and the looping index to zero
		iLay  = 0
		iLoop = 0

		# Perform a while loop over the layers
		while iLay < numberOfLayers:

			# Set the receiver index equal to the looping index
			iRec = iLoop

			# Perform a while loop over the receivers, starting at the current index of the receivers
			while iRec < numberOfReceivers:

				# Check if the depth of the receiver is smaller than or equal to the layer depth
				if receiverVector[ iRec ] <= layerVector[ iLay ]:

					# Check if this is the first layer
					if iLay == 0:

						# Set the depth equal to the receiver depth
						depths[ iRec, iLay ] = receiverVector[ iRec ]
						
					# Check if this is not the first layer
					else:

						# Set the depth equal to the difference between the receiver and layer
						depths[ iRec, iLay ] = receiverVector[ iRec ] - layerVector[ iLay - 1 ]

					# Set the looping index equal to the receiver index  + 1
					iLoop = iRec + 1

				# Check if the depth of the receiver is greater than the layer depth
				else:

					# Check if this is the first layer
					if iLay == 0:

						# Set the depth equal to the layer depth
						depths[ iRec, iLay ] = layerVector[ iLay ]
						
					# Check if this is not the first layer
					else:

						# Set the depth equal to the difference between the layers
						depths[ iRec, iLay ] = layerVector[ iLay ] - layerVector[ iLay - 1 ]

				# Increment the receiver index by 1
				iRec += 1

			# Increment the layer index by 1
			iLay += 1

	# Return the depth matrix
	return( depths )

# Create a function for picking the 1D VSP
def pickVSPzo( receiverVector, fig, ax ):

	# Determine the receiver spacing
	receiverSpacing = receiverVector[ 1 ] - receiverVector[ 0 ]

	# Create empty lists that hold the values for the picked times for each receiver
	clickedT = []
	clickedR = []

	# Create an array of travelTimes
	travelTimes = np.zeros_like( receiverVector )

	# Plot the picks for the receiver locations
	scatter_plot = ax.scatter(travelTimes, receiverVector, c='green', s=40)

	# Click event handler
	def onclick(event):
		if event.inaxes == ax and event.xdata is not None and event.ydata is not None:
			
			# Get the data
			tPick, rPick = event.xdata, event.ydata

			# Snap to the nearest receiver index within bounds
			iRec = np.clip( int( np.round( ( rPick - receiverVector.min() ) / receiverSpacing ) ), 0, len( receiverVector ) - 1 )
			rPick = iRec * receiverSpacing + receiverVector.min()

			# Check if the pick for this receiver is already in the list
			if rPick in clickedR:

				# Change the value to the new pick
				clickedR[ clickedR.index( rPick ) ] = tPick

			# Otherwise add to list of picks
			else: 

				# Add to the list
				clickedT.append( tPick )
				clickedR.append( rPick )

			# Sort the lists so they are purely ascending
			rval, tval = zip( *sorted( zip( clickedR,  clickedT ) ) )

			# Interpolate based on the picks
			travelTimes[ : ] = np.interp( receiverVector, rval, tval )
			
			# Update scatter plot
			scatter_plot.set_offsets(list(zip(travelTimes, receiverVector)))
			fig.canvas.draw()

	# Connect the click event to handler
	cid = fig.canvas.mpl_connect('button_press_event', onclick)

	plt.show( block = True )

	return ( travelTimes )