import numpy as np
from scipy.interpolate import griddata


def bin_distances(distances, bin_size=10):
    # Bin the distances into groups of `bin_size` kilometers
    binned_distances = {}
    for i, distance in enumerate(distances):
        bin_index = distance // bin_size
        if bin_index not in binned_distances:
            binned_distances[bin_index] = (distance, i)
        elif i < binned_distances[bin_index][1]:
            binned_distances[bin_index] = (distance, i)

    # Select the first distance in each bin and its index
    first_distances = []
    for bin_index in binned_distances:
        first_distance, first_distance_index = binned_distances[bin_index]
        first_distances.append(first_distance_index)

    return first_distances


def interpolate_vel_model(
    velocity_model,
    initial_velocity,
    lat_values,
    lon_values,
    depth_values,
    n_lat,
    n_lon,
    n_depth,
):
    # Create a mask for points with the initial velocity
    initial_velocity_mask = velocity_model == initial_velocity

    # Find the indices of points with non-initial velocities
    non_initial_velocity_indices = np.argwhere(~initial_velocity_mask)

    # Extract the coordinates and corresponding velocities of the known points
    known_points = np.column_stack(
        [
            lat_values[non_initial_velocity_indices[:, 0]],
            lon_values[non_initial_velocity_indices[:, 1]],
            depth_values[non_initial_velocity_indices[:, 2]],
        ]
    )

    # Find the maximum depth in the known_points
    max_known_depth = np.max(known_points[:, 2])

    known_velocities = velocity_model[~initial_velocity_mask]

    # Create a grid of points for the entire volume
    grid_points = (
        np.array(np.meshgrid(lat_values, lon_values, depth_values, indexing="ij"))
        .reshape(3, -1)
        .T
    )

    # Create a mask for grid points that are deeper than the maximum known depth
    depth_mask = grid_points[:, 2] <= max_known_depth

    # Interpolate the velocities at the grid points
    interpolated_velocities = griddata(
        known_points, known_velocities, grid_points[depth_mask], method="linear"
    )

    # Fill nan values with the nearest known velocities
    interpolated_velocities_filled = griddata(
        known_points, known_velocities, grid_points[depth_mask], method="nearest"
    )
    interpolated_velocities[
        np.isnan(interpolated_velocities)
    ] = interpolated_velocities_filled[np.isnan(interpolated_velocities)]

    # Initialize an array with the same length as grid_points and fill it with nan values
    interpolated_velocities_with_depth_limit = np.full(grid_points.shape[0], np.nan)

    # Update the array with the interpolated velocities for the masked grid points
    interpolated_velocities_with_depth_limit[depth_mask] = interpolated_velocities

    # Reshape the interpolated velocities to match the shape of the velocity_model
    interpolated_velocity_model = interpolated_velocities_with_depth_limit.reshape(
        n_lat, n_lon, n_depth
    )

    return interpolated_velocity_model


# Function to find the closest index for a given value in an array
def find_closest_index(array, value):
    return np.argmin(np.abs(array - value))