def assign_jobs_to_workers(gdf_homes):
    """
    Assign job locations to workers residing in the buildings of gdf_homes.

    For each building in the GeoDataFrame, the function takes the potential job locations and 
    assigns these to every worker residing in that building.

    Parameters:
    - gdf_homes (GeoDataFrame): A GeoDataFrame containing home data. It should have 'building' 
      and 'potential_jobs' columns, where 'building' objects have a 'worker_list' attribute, and each 
      worker in this list has a 'work_location' attribute.

    Returns:
    - None: The function modifies the 'work_location' attribute of each worker in-place.
    """

    # Go through each row of the GeoDataFrame
    for _, row in gdf_homes.iterrows():
        building = row['building']           # Retrieve the building object
        job_locations = row['potential_jobs'] # Retrieve the potential job locations for this building

        # Iterate over each worker in the building
        for worker,job_point in zip(building.worker_list, job_locations):
            # Assign the job location to the worker's work_location attribute
            job_location = Location("WORK", "Work", job_point, "WORK")
            worker.work_location = job_location

def assign_workers_to_buildings():
    """ 
    Assigns workers to buildings based on their primary status.
    """
    # Add workers to buildings
    for building in Building.instances:
        houses = building.houses
        for house in houses:
            household = house.household
            for member in household.members:
                # Workers is a list of workers with primary_status = "WORK"
                if member.primary_status == "WORK":
                    building.workers += 1
                    building.worker_list.append(member)

def compute_job_density(file_path, radius=1000):
    """
    Computes job densities for each point in the GeoDataFrame using the given radius.

    The function loads a shapefile, calculates the centroid for each geometry and then computes
    the job density around each centroid using a given radius. The BallTree data structure is 
    utilized for efficient spatial queries.

    Parameters:
    - file_path (str): Path to the shapefile.
    - radius (float): Radius in meters for which job density is calculated. Default is 1000 meters.

    Returns:
    - gdf_jobs (GeoDataFrame): Updated GeoDataFrame with job densities.
    """
    logger.info("Computing job densities...")
    # Load data from the shapefile
    gdf = gpd.read_file(file_path)

    # Compute centroids for each geometry in the GeoDataFrame
    gdf['centroid'] = gdf['geometry'].centroid

    # Create a new GeoDataFrame with only the 'jobs' column and the computed centroids
    # Set the 'centroid' column as the geometry for this new GeoDataFrame
    gdf_jobs = gdf[['jobs', 'centroid']].copy()
    gdf_jobs = gdf_jobs.set_geometry('centroid')

    # Convert the GeoDataFrame's point geometries to a matrix format suitable for BallTree
    coordinates = np.array(list(zip(gdf_jobs.geometry.x, gdf_jobs.geometry.y)))

    # Construct a BallTree for efficient spatial queries
    ball_tree = BallTree(coordinates)

    # Define an inner function to perform spatial queries and calculate job density
    def job_density(point):
        """
        Compute job density for a given point by querying nearby jobs within the specified radius.

        Parameters:
        - point (tuple): A tuple representing the x and y coordinates of the point.

        Returns:
        - float: Job density for the given point.
        """
        # Get the indices of points within the given radius using the BallTree
        indices = ball_tree.query_radius([point], r=radius, return_distance=False)[0]
        # Sum up the jobs for the queried points to compute the density
        return gdf_jobs.iloc[indices]['jobs'].sum()

    # Compute job densities for each point in the GeoDataFrame
    gdf_jobs['job_density'] = [job_density(point) for point in coordinates]

    return gdf_jobs

def compute_preferred_locations(add_proposed_locations = False):
    """
    #TODO This can be made faster by grouping buildings into chunks of 200x200m and assigning them together.
    Computes and assigns preferred locations for each building instance based on their coordinates.

    For each building in the list of instances, this function:
    1. Determines the building's coordinate.
    2. Finds the closest general locations to that coordinate.
    3. Finds the closest grocery locations to that coordinate.
    4. Merges the two lists of locations.
    5. Initializes a PreferredLocations object using the combined locations.
    6. Assigns the PreferredLocations object to the building's preferred_locations attribute.
    """
    logger.info("Computing preferred locations...")
    gdf_amenities = gpd.read_file(ALL_AMENITIES_PATH)

    if add_proposed_locations:
        # Check if the proposed amenities file exists
        if os.path.exists(PROPOSED_AMENITIES_PATH):

            # Load the proposed amenities GeoJSON file
            gdf_proposed_amenities = gpd.read_file(PROPOSED_AMENITIES_PATH)

            if len(gdf_proposed_amenities) > 0:

                logger.info("Proposed amenities file found. Combining amenities for location assignment")
                # Set the CRS for the proposed amenities GeoDataFrame
                gdf_proposed_amenities.set_crs(epsg=3006, inplace=True,allow_override=True)

                # Concatenate the existing and proposed amenities into one GeoDataFrame
                gdf_amenities = pd.concat([gdf_amenities, gdf_proposed_amenities], ignore_index=True)

                logger.info(gdf_proposed_amenities.activity.value_counts())
        else:
            logger.info("Proposed amenities are empty or no proposed amenities file found. Proceeding with location assignment")

    # Initialize the LocationFinder object outside the loop 
    # (since it probably doesn't need to be re-initialized for each building).
    location_finder = LocationFinder(gdf_amenities)

    # Loop through all the building instances
    for building in Building.instances:
        # Extract the coordinate (centroid) of the building
        origin = building.coord

        # Find the closest general locations to the building's coordinate
        locations = location_finder.find_closest_locations(origin)

        # Find the closest grocery locations to the building's coordinate
        grocery_locations = location_finder.find_closest_grocery_locations(origin)

        # Initialize a PreferredLocations object using the combined lists of locations
        preferred_locations = PreferredLocations(locations + grocery_locations)
        preferred_locations.origin = origin

        # Assign the generated PreferredLocations object to the building's attribute
        building.preferred_locations = preferred_locations

    logger.info(f"Preferred locations computed for {len(Building.instances)} buildings.")

def gravity_model_simulation(gdf_homes, gdf_jobs, density_weight=0.5,  distance_decay=2.5, plot=True):
    """
    Simulate job attraction based on the gravity model.

    The gravity model is a spatial interaction model that suggests that interaction between two places
    (for example, the number of people that commute from one place to another for work) is proportional to 
    the product of the population of the two places and inversely proportional to the square of the distance 
    between them.

    Parameters:
    - gdf_homes (GeoDataFrame): A GeoDataFrame containing home data. It must have 'footprint' and 'workers' columns.
    - gdf_jobs (GeoDataFrame): A GeoDataFrame containing job location data. It must have a 'job_density' column.
    - density_weight (float): The weight to apply to job density in the attraction calculation. Default is 1.
    - distance_decay (float): The power to which distance is raised in the attraction calculation. Default is 2.
    - plot (bool): Whether or not to plot the results. Default is True.

    Returns:
    - gdf_worker_jobs (GeoDataFrame): A GeoDataFrame representing the relationship between homes and potential job locations.

    To skew the attraction towards job density - increase the density_weight.
    To skew the attraction towards distance - increase the distance_decay.

    """

    # 1. Calculate Distance Matrix
    # Extracting the coordinates of home and job centroids
    homes_coords = gdf_homes.footprint.centroid.apply(lambda geom: (geom.x, geom.y)).tolist()
    jobs_coords = gdf_jobs.centroid.apply(lambda geom: (geom.x, geom.y)).tolist()
    # Computing the distance matrix
    distance_matrix = pairwise_distances(homes_coords, jobs_coords, metric="euclidean")

    # 2. Gravity Model Calculation
    # Compute attraction based on job density and distance decay
    attraction_matrix = (gdf_jobs['job_density'].values ** density_weight) / (distance_matrix ** distance_decay)

    def get_job_locations(row_index, n_jobs):
        """
        For each home, get the most attractive job locations based on the gravity model.

        Parameters:
        - row_index (int): The index of the home row.
        - n_jobs (int): The number of jobs associated with the home.

        Returns:
        - list: A list of the most attractive job locations' centroids.
        """
        # Getting the indices of job locations with the highest attraction
        top_jobs_indices = np.argsort(attraction_matrix[row_index])[-n_jobs:]
        unique_jobs = set()
        # Picking the top n_jobs from the list
        for idx in reversed(top_jobs_indices):
            if len(unique_jobs) < n_jobs:
                unique_jobs.add(idx)
        return gdf_jobs.iloc[list(unique_jobs)]['centroid'].tolist()

    # Assign potential job locations for each home based on number of workers and attraction
    gdf_homes['potential_jobs'] = [get_job_locations(idx, int(workers)) for idx, workers in enumerate(gdf_homes['workers'])]

    logger.info("Assigning jobs to workers...")

    # Assign job locations to workers
    assign_jobs_to_workers(gdf_homes)

    # Data aggregation
    # Constructing the resultant data
    data = []
    for _, home_row in gdf_homes.iterrows():
        home_footprint = home_row.footprint
        home_centroid = home_row.footprint.centroid
        for job_point in home_row['potential_jobs']:
            data.append({
                'home_footprint': home_footprint,
                'home_centroid': home_centroid,
                'worker': 1,
                'job_location': job_point
            })

    # Convert the list of dictionaries to a DataFrame and then to a GeoDataFrame
    df = pd.DataFrame(data)
    gdf_worker_jobs = gpd.GeoDataFrame(df, geometry='job_location', crs=gdf_homes.crs)

    if plot:
        # Visualize the result
        fig, ax = plt.subplots(figsize=(10,10))
        gpd.GeoSeries(gdf_worker_jobs['home_footprint']).plot(ax=ax, color='grey', alpha=0.5)
        gpd.GeoSeries(gdf_worker_jobs['home_centroid']).plot(ax=ax, color='red', markersize=1)
        gpd.GeoSeries(gdf_worker_jobs['job_location']).plot(ax=ax, color='blue', markersize=1)
        plt.show()

    return gdf_worker_jobs