From 98b63541652880aaeb7fec2dc3d169ccb5dcfa79 Mon Sep 17 00:00:00 2001
From: "WQY\\qiong" <qiongyu.wang@outlook.com>
Date: Sun, 30 Apr 2023 23:58:22 +0800
Subject: [PATCH] Support dma calculation

---
 api/__init__.py                   |  3 ++
 api/s35_district_metering_area.py | 85 +++++++++++++++++++++++++++++++
 tjnetwork.py                      |  7 +++
 3 files changed, 95 insertions(+)
 create mode 100644 api/s35_district_metering_area.py

diff --git a/api/__init__.py b/api/__init__.py
index bf65fe5..4d98d94 100644
--- a/api/__init__.py
+++ b/api/__init__.py
@@ -139,6 +139,9 @@ 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 .s35_district_metering_area import PARTITION_TYPE_RB, PARTITION_TYPE_KWAY
+from .s35_district_metering_area import calculate_district_metering_area
+
 from .s36_service_area import calculate_service_area
 
 from .s37_virtual_district import calculate_virtual_district
diff --git a/api/s35_district_metering_area.py b/api/s35_district_metering_area.py
new file mode 100644
index 0000000..d907374
--- /dev/null
+++ b/api/s35_district_metering_area.py
@@ -0,0 +1,85 @@
+import ctypes
+import os
+from .database import *
+from .s32_region_util import Topology
+
+
+PARTITION_TYPE_RB = 0
+PARTITION_TYPE_KWAY = 1
+
+
+def calculate_district_metering_area(name: str, nodes: list[str], part_count: int = 1, part_type: int = PARTITION_TYPE_RB) -> list[list[str]]:
+    if part_type != PARTITION_TYPE_RB and part_type != PARTITION_TYPE_KWAY:
+        return []
+    if part_count <= 0:
+        return []
+    elif part_count == 1:
+        return [nodes]
+
+    lib = ctypes.CDLL(os.path.join(os.getcwd(), 'api', 'CMetis.dll'))
+
+    METIS_NOPTIONS = 40
+    c_options = (ctypes.c_int64 * METIS_NOPTIONS)()
+
+    METIS_OK = 1
+    result = lib.set_default_options(c_options)
+    if result != METIS_OK:
+        return []
+
+    METIS_OPTION_PTYPE , METIS_OPTION_CONTIG = 0, 13
+    c_options[METIS_OPTION_PTYPE] = part_type
+    c_options[METIS_OPTION_CONTIG] = 1
+
+    topology = Topology(name, nodes)
+    t_nodes = topology.nodes()
+    t_links = topology.links()
+    t_node_list = topology.node_list()
+    t_link_list = topology.link_list()
+
+    nedges = len(t_link_list) * 2
+
+    c_nvtxs = ctypes.c_int64(len(t_node_list))
+    c_ncon = ctypes.c_int64(1)
+    c_xadj = (ctypes.c_int64 * (c_nvtxs.value + 1))()
+    c_adjncy = (ctypes.c_int64 * nedges)()
+    c_vwgt = (ctypes.c_int64 * (c_ncon.value * c_nvtxs.value))()
+    c_adjwgt = (ctypes.c_int64 * nedges)()
+    c_vsize = (ctypes.c_int64 * c_nvtxs.value)()
+
+    c_xadj[0] = 0
+
+    l, n = 0, 0
+    c_xadj_i = 1
+    for node in t_node_list:
+        links = t_nodes[node]['links']
+        for link in links:
+            node1 = t_links[link]['node1']
+            node2 = t_links[link]['node2']
+            c_adjncy[l] = t_node_list.index(node2) if node2 != node else t_node_list.index(node1)
+            c_adjwgt[l] = 1
+            l += 1
+        if len(links) > 0:
+            c_xadj[c_xadj_i] = l # adjncy.size()
+            c_xadj_i += 1
+        c_vwgt[n] = 1
+        c_vsize[n] = 1
+        n += 1
+
+    part_func = lib.part_graph_recursive if part_type == PARTITION_TYPE_RB else lib.part_graph_kway
+
+    c_nparts = ctypes.c_int64(part_count)
+    c_tpwgts = ctypes.POINTER(ctypes.c_double)()
+    c_ubvec = ctypes.POINTER(ctypes.c_double)()
+    c_out_edgecut = ctypes.c_int64(0)
+    c_out_part = (ctypes.c_int64 * c_nvtxs.value)()
+    result = part_func(ctypes.byref(c_nvtxs), ctypes.byref(c_ncon), c_xadj, c_adjncy, c_vwgt, c_vsize, c_adjwgt, ctypes.byref(c_nparts), c_tpwgts, c_ubvec, c_options, ctypes.byref(c_out_edgecut), c_out_part)
+    if result != METIS_OK:
+        return []
+
+    dmas : list[list[str]]= []
+    for i in range(part_count):
+        dmas.append([])
+    for i in range(c_nvtxs.value):
+        dmas[c_out_part[i]].append(t_node_list[i])
+
+    return dmas
diff --git a/tjnetwork.py b/tjnetwork.py
index a30ee82..522f931 100644
--- a/tjnetwork.py
+++ b/tjnetwork.py
@@ -167,6 +167,10 @@ DISTRIBUTION_TYPE_ADD      = api.DISTRIBUTION_TYPE_ADD
 DISTRIBUTION_TYPE_OVERRIDE = api.DISTRIBUTION_TYPE_OVERRIDE
 
 
+PARTITION_TYPE_RB   = api.PARTITION_TYPE_RB
+PARTITION_TYPE_KWAY = api.PARTITION_TYPE_KWAY
+
+
 ############################################################
 # project
 ############################################################
@@ -998,6 +1002,9 @@ def distribute_demand_to_region(name: str, demand: float, region: str, type: str
 # district_metering_area 35
 ############################################################
 
+def calculate_district_metering_area(name: str, nodes: list[str], part_count: int = 1, part_type: int = PARTITION_TYPE_RB) -> list[list[str]]:
+    return api.calculate_district_metering_area(name, nodes, part_count, part_type)
+
 
 ############################################################
 # service_area 36