import itertools

from qgis.PyQt.QtCore import QVariant
from qgis.core import (QgsFeature, QgsField, QgsGeometry, QgsPointXY, 
                       QgsVectorLayer, QgsProject, QgsVectorDataProvider)

print("Creating nodes...")

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

# Prepare the NODES layer
nodes_layer = QgsVectorLayer("Point", "NODES", "memory")
prov = nodes_layer.dataProvider()

# Add 'id' field to the NODES layer
prov.addAttributes([QgsField("id", QVariant.Int)])
nodes_layer.updateFields()

# Prepare the start_node and end_node fields for the EDGES layer
edges_layer.startEditing()
edges_prov = edges_layer.dataProvider()
edges_prov.addAttributes([QgsField("start_node", QVariant.Int),
                          QgsField("end_node", QVariant.Int)])
edges_layer.updateFields()

# To store the unique nodes
unique_nodes = []

id_counter = itertools.count(start=1)

node_id_map = {}

total_features = edges_layer.featureCount()
count = 0

for feature in edges_layer.getFeatures():
    count += 1
    print(f'Processing feature {count} of {total_features} ({(count/total_features)*100:.2f}%)')
    
    # Get the start and end points
    start_point = feature.geometry().asPolyline()[0]
    end_point = feature.geometry().asPolyline()[-1]
    
    # Check if nodes are unique and add them to NODES layer
    if start_point not in unique_nodes:
        unique_nodes.append(start_point)
        node_id = next(id_counter)
        node_feature = QgsFeature()
        node_feature.setGeometry(QgsGeometry.fromPointXY(start_point))
        node_feature.setAttributes([node_id])
        prov.addFeature(node_feature)
        node_id_map[start_point] = node_id
        
    if end_point not in unique_nodes:
        unique_nodes.append(end_point)
        node_id = next(id_counter)
        node_feature = QgsFeature()
        node_feature.setGeometry(QgsGeometry.fromPointXY(end_point))
        node_feature.setAttributes([node_id])
        prov.addFeature(node_feature)
        node_id_map[end_point] = node_id

    # Populate the start_node and end_node fields for the EDGES layer
    edges_layer.changeAttributeValue(feature.id(), 
                                     edges_layer.fields().indexOf("start_node"), 
                                     node_id_map[start_point])
    edges_layer.changeAttributeValue(feature.id(), 
                                     edges_layer.fields().indexOf("end_node"), 
                                     node_id_map[end_point])

# Stop editing the EDGES layer and save changes
edges_layer.commitChanges()

# Add the layer to the Layers panel
QgsProject.instance().addMapLayer(nodes_layer)

print("Done.")