# Import required libraries
import networkx as nx
import time
from qgis.PyQt.QtCore import QVariant
from qgis.core import QgsProject, QgsVectorLayer, QgsFeature, QgsField, QgsGeometry, QgsPoint
import heapq

# Get layers
edges_layer = QgsProject.instance().mapLayersByName('EDGES')[0]
nodes_layer = QgsProject.instance().mapLayersByName('NODES')[0]
home_points_layer = QgsProject.instance().mapLayersByName('HOME_POINTS')[0]

# Build the graph
G = nx.Graph()
nodes_data = {feature['id']: feature.geometry() for feature in nodes_layer.getFeatures()}  # store node geometries by id
edges_data = {(feature['start_node'], feature['end_node']): (feature.geometry(), feature['length'], feature['cost']) for feature in edges_layer.getFeatures()}  # store edge geometries by node pairs

for feature in nodes_layer.getFeatures():
    G.add_node(feature['id'])

for feature in edges_layer.getFeatures():
    start_node = feature['start_node']  # assuming edges layer has start_node & end_node fields
    end_node = feature['end_node']
    G.add_edge(start_node, end_node, weight=feature['cost'], type=feature['type'], length=feature['length'], cost=feature['cost'])

# Identify home nodes
home_nodes = {feature['drop_point_id'] for feature in home_points_layer.getFeatures() if feature['drop_point_id'] in G}

# Create a subgraph of G that only includes the home nodes and any nodes that connect them
subgraph = nx.Graph()
visited = set()

# Start at any 'home' node
start = next(iter(home_nodes))
visited.add(start)

counter = 0
total = len(home_nodes)
process_times = []

while home_nodes.difference(visited):
    start_time = time.time()

    # Find the nearest 'home' node that is not in our tree yet
    next_node = min(
        (node for node in home_nodes if node not in visited),
        key=lambda node: nx.shortest_path_length(G, start, node, weight='cost')
    )

    # Create a subgraph excluding other home nodes except start and next_node
    G_sub = G.copy()
    remove_nodes = home_nodes - {start} - {next_node}
    G_sub.remove_nodes_from(remove_nodes)

    # Add the shortest path to this node to our tree
    shortest_path = nx.shortest_path(G_sub, start, next_node, weight='cost')
    subgraph.add_edges_from(zip(shortest_path, shortest_path[1:]))
    visited.add(next_node)
    start = next_node

    end_time = time.time()
    process_times.append(end_time - start_time)
    counter += 1
    avg_time = sum(process_times) / len(process_times)
    remaining = (total - counter) * avg_time
    hours, remainder = divmod(remaining, 3600)
    minutes, seconds = divmod(remainder, 60)
    print(f'Processing home node {counter} of {total}.\nEstimated time remaining: {int(hours)} hours, {int(minutes)} minutes, and {int(seconds)} seconds')
    
# Create a minimum spanning tree from the subgraph
mst = nx.minimum_spanning_tree(subgraph, weight='cost')

# Prepare the NETWORK layer
network_layer = QgsVectorLayer("LineString", "NETWORK", "memory")
network_prov = network_layer.dataProvider()

# Add 'id', 'type', 'length', and 'cost' fields to the NETWORK layer
network_prov.addAttributes([QgsField("id", QVariant.Int), QgsField("type", QVariant.String), QgsField("length", QVariant.Double), QgsField("cost", QVariant.Double)])
network_layer.updateFields()

# Add edges from the minimum spanning tree to the NETWORK layer
for edge in mst.edges():
    edge_geometry, edge_length, edge_cost = edges_data[(edge[0], edge[1])] if (edge[0], edge[1]) in edges_data else edges_data[(edge[1], edge[0])]  # fetch geometry, length, cost by node pair
    edge_type = G.get_edge_data(edge[0], edge[1])['type']  # get edge 'type' from the original graph G
    network_feature = QgsFeature()
    network_feature.setGeometry(edge_geometry)
    network_feature.setAttributes([edge[0], edge_type, edge_length, edge_cost])  # set the type, length, cost attribute from edge data
    network_prov.addFeature(network_feature)

# Update NETWORK layer
network_layer.updateExtents()

# Add NETWORK layer to Layers panel
QgsProject.instance().addMapLayer(network_layer)
print("Done.")