from geopy.distance import great_circle
from qgis.core import (QgsFeature, QgsField, QgsGeometry, QgsPoint, QgsPointXY, QgsVectorLayer, QgsProject, QgsSpatialIndex)

# Locate the aerial path layer
aerial_path_layer = QgsProject.instance().mapLayersByName('AERIAL_PATH')[0]

# Locate the edges layer
edges_layer = QgsProject.instance().mapLayersByName('EDGES')[0]

# Define the data provider
pr = edges_layer.dataProvider()

# Initialize spatial index
index = QgsSpatialIndex()

# List to hold new features and their IDs
new_edges = []
features_by_id = {}

# Iterate over every feature of the AERIAL_PATH layer
for feature in aerial_path_layer.getFeatures():
    # Get the geometry of the feature
    geom = feature.geometry()

    # Handle both 2D and 3D geometries
    if geom.isMultipart():
        vertices = [v for part in geom.asMultiPolyline() for v in part]
    else:
        vertices = geom.asPolyline()

    for i in range(1, len(vertices)):
        # Calculate the distance between the two points in feet
        pt1 = vertices[i - 1]
        pt2 = vertices[i]
        length = great_circle((pt1.y(), pt1.x()), (pt2.y(), pt2.x())).feet

        # Create a new feature in the EDGES layer for each vertex pair
        edge = QgsFeature()
        new_geom = QgsGeometry.fromPolyline([QgsPoint(pt1), QgsPoint(pt2)])
        edge.setGeometry(new_geom)
        edge.setAttributes(['Aerial', length, length * 2.5])

        new_edges.append(edge)
        pr.addFeature(edge)
        index.insertFeature(edge)
        features_by_id[edge.id()] = edge

# Start editing the edges layer
edges_layer.startEditing()

# Adjust intersecting line endpoints to the average intersection point
for edge_id, edge in features_by_id.items():
    endpoints = [QgsPoint(point) for point in edge.geometry().asPolyline()]
    for i in [0, -1]:  # Check both endpoints
        # Find nearby endpoints (within a small threshold)
        search_radius = 0.00005
        nearby_ids = index.intersects(QgsGeometry.fromPointXY(QgsPointXY(endpoints[i])).buffer(search_radius, 8).boundingBox())
        nearby_endpoints = [QgsPoint(point) for id in nearby_ids for point in features_by_id[id].geometry().asPolyline() if QgsPoint(point).distance(endpoints[i]) < search_radius]
        if len(nearby_endpoints) > 2:
            # Update endpoint to the average nearby endpoint
            avg_x = sum(point.x() for point in nearby_endpoints) / len(nearby_endpoints)
            avg_y = sum(point.y() for point in nearby_endpoints) / len(nearby_endpoints)
            endpoints[i] = QgsPoint(avg_x, avg_y)  # Create new QgsPoint object
    # Update geometry in the layer
    edges_layer.changeGeometry(edge.id(), QgsGeometry.fromPolyline(endpoints))

# Commit the changes and update the layer's extent when new features have been added
edges_layer.commitChanges()
edges_layer.updateExtents()
print("Done.")