From 7a1a8a3a7bfa63874a0df1332e30eb39e6606416 Mon Sep 17 00:00:00 2001
From: "WQY\\qiong" <qiongyu.wang@outlook.com>
Date: Sun, 30 Apr 2023 18:43:44 +0800
Subject: [PATCH] Support service area calculation

---
 api/__init__.py         |   2 +
 api/s36_service_area.py | 119 ++++++++++++++++++++++++++++++++++++++++
 tjnetwork.py            |   3 +
 3 files changed, 124 insertions(+)
 create mode 100644 api/s36_service_area.py

diff --git a/api/__init__.py b/api/__init__.py
index c22bd76..bf65fe5 100644
--- a/api/__init__.py
+++ b/api/__init__.py
@@ -139,4 +139,6 @@ from .s33_region import get_region_schema, get_region, set_region, add_region, d
 from .s34_water_distribution import DISTRIBUTION_TYPE_ADD, DISTRIBUTION_TYPE_OVERRIDE
 from .s34_water_distribution import distribute_demand_to_nodes, distribute_demand_to_region
 
+from .s36_service_area import calculate_service_area
+
 from .s37_virtual_district import calculate_virtual_district
diff --git a/api/s36_service_area.py b/api/s36_service_area.py
new file mode 100644
index 0000000..ec012e5
--- /dev/null
+++ b/api/s36_service_area.py
@@ -0,0 +1,119 @@
+import sys
+import json
+from queue import Queue
+from .database import *
+from .s0_base import get_node_links, get_link_nodes
+
+sys.path.append('..')
+from epanet.epanet import run_project
+
+def calculate_service_area(name: str, time_index: int = 0) -> dict[str, Any]:
+    inp = json.loads(run_project(name))
+
+    time_count = len(inp['node_results'][0]['result'])
+    if time_index >= time_count:
+        return {}
+
+    sources : dict[str, list[str]] = {}
+    for node_result in inp['node_results']:
+        result = node_result['result'][time_index]
+        if result['demand'] < 0:
+            sources[node_result['node']] = []
+
+    link_flows: dict[str, float] = {}
+    for link_result in inp['link_results']:
+        result = link_result['result'][time_index]
+        link_flows[link_result['link']] = float(result['flow'])
+
+    # build source to nodes map
+    for source in sources:
+        queue = Queue()
+        queue.put(source)
+
+        while not queue.empty():
+            cursor = queue.get()
+            if cursor not in sources[source]:
+                sources[source].append(cursor)
+
+                links = get_node_links(name, cursor)
+                for link in links:
+                    node1, node2 = get_link_nodes(name, link)
+                    if node1 == cursor and link_flows[link] > 0:
+                        queue.put(node2)
+                    elif node2 == cursor and link_flows[link] < 0:
+                        queue.put(node1)
+
+    # calculation concentration
+    concentration_map: dict[str, dict[str, float]] = {}
+    node_wip: list[str] = []
+    for source, nodes in sources.items():
+        for node in nodes:
+            if node not in concentration_map:
+                concentration_map[node] = {}
+            concentration_map[node][source] = 0.0
+            if node not in node_wip:
+                node_wip.append(node)
+
+    # if only one source, done
+    for node, concentrations in concentration_map.items():
+        if len(concentrations) == 1:
+            node_wip.remove(node)
+            for key in concentrations.keys():
+                concentration_map[node][key] = 1.0
+
+    node_upstream : dict[str, list[tuple[str, str]]] = {}
+    for node in node_wip:
+        if node not in node_upstream:
+            node_upstream[node] = []
+
+        links = get_node_links(name, node)
+        for link in links:
+            node1, node2 = get_link_nodes(name, link)
+            if node2 == node and link_flows[link] > 0:
+                node_upstream[node].append((link, node1))
+            elif node1 == node and link_flows[link] < 0:
+                node_upstream[node].append((link, node2))
+
+    while len(node_wip) != 0:
+        done = []
+        for node in node_wip:
+            up_link_nodes = node_upstream[node]
+            ready = True
+            for link_node in up_link_nodes:
+                if link_node in node_wip:
+                    ready = False
+                    break
+            if ready:
+                for link_node in up_link_nodes:
+                    for source, concentration in concentration_map[link_node[1]].items():
+                        concentration_map[node][source] += concentration * abs(link_flows[link_node[0]])
+
+                # normalize
+                sum = 0.0
+                for source, concentration in concentration_map[node].items():
+                    sum += concentration
+                for source in concentration_map[node].keys():
+                    concentration_map[node][source] /= sum
+
+                done.append(node)
+
+        for node in done:
+            node_wip.remove(node)
+
+    source_to_main_node: dict[str, list[str]] = {}
+    for node, value in concentration_map.items():
+        max_source = ''
+        max_concentration = 0.0
+        for s, c in value.items():
+            if c > max_concentration:
+                max_concentration = c
+                max_source = s
+        if max_source not in source_to_main_node:
+            source_to_main_node[max_source] = []
+        source_to_main_node[max_source].append(node)
+
+    sas: list[dict[str, Any]] = []
+    for source, nodes in source_to_main_node.items():
+        sas.append({ 'source': source, 'nodes': nodes })
+
+    return { 'service_areas' : sas, 'concentrations': concentration_map }
diff --git a/tjnetwork.py b/tjnetwork.py
index 0517260..a30ee82 100644
--- a/tjnetwork.py
+++ b/tjnetwork.py
@@ -1003,6 +1003,9 @@ def distribute_demand_to_region(name: str, demand: float, region: str, type: str
 # service_area 36
 ############################################################
 
+def calculate_service_area(name: str, time_index: int = 0) -> dict[str, Any]:
+    return api.calculate_service_area(name, time_index)
+
 
 ############################################################
 # virtual_district 37