import ctypes
import os
import numpy as np
import pymetis
from .database import *
from .s0_base import get_nodes
from .s32_region_util import get_nodes_in_region
from .s32_region_util import Topology


PARTITION_TYPE_RB = 0
PARTITION_TYPE_KWAY = 1

'''
adjacency_list = [np.array([4, 2, 1]),
                  np.array([0, 2, 3]),
                  np.array([4, 3, 1, 0]),
                  np.array([1, 2, 5, 6]),
                  np.array([0, 2, 5]),
                  np.array([4, 3, 6]),
                  np.array([5, 3])]
n_cuts, membership = pymetis.part_graph(2, adjacency=adjacency_list)
# n_cuts = 3
# membership = [1, 1, 1, 0, 1, 0, 0]

nodes_part_0 = np.argwhere(np.array(membership) == 0).ravel() # [3, 5, 6]
nodes_part_1 = np.argwhere(np.array(membership) == 1).ravel() # [0, 1, 2, 4]

print(nodes_part_0)
print(nodes_part_1)
'''


def calculate_district_metering_area_for_nodes(name: str, nodes: list[str], part_count: int = 1, part_type: int = PARTITION_TYPE_RB) -> list[list[str]]:
    topology = Topology(name, nodes)
    t_nodes = topology.nodes()
    t_links = topology.links()
    t_node_list = topology.node_list()

    adjacency_list = []

    for node in t_node_list:
        links: list[str] = t_nodes[node]['links']
        a_nodes: list[int] = []
        for link in links:
            if t_links[link]['node1'] == node:
                i = t_node_list.index(t_links[link]['node2'])
                a_nodes.append(i)
            elif t_links[link]['node2'] == node:
                i = t_node_list.index(t_links[link]['node1'])
                a_nodes.append(i)
        adjacency_list.append(np.array(a_nodes))

    recursive = part_type == PARTITION_TYPE_RB
    n_cuts, membership = pymetis.part_graph(nparts=part_count, adjacency=adjacency_list, recursive=recursive, contiguous=True)

    result: list[list[str]] = []
    for i in range(0, part_count):
        indices: list[int] = list(np.argwhere(np.array(membership) == i).ravel())
        index_strs: list[str] = []
        for index in indices:
            index_strs.append(t_node_list[index])
        result.append(index_strs)

    return result


def _calculate_district_metering_area_for_nodes(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


def calculate_district_metering_area_for_region(name: str, region: str, part_count: int = 1, part_type: int = PARTITION_TYPE_RB) -> list[list[str]]:
    nodes = get_nodes_in_region(name, region)
    return calculate_district_metering_area_for_nodes(name, nodes, part_count, part_type)


def calculate_district_metering_area_for_network(name: str, part_count: int = 1, part_type: int = PARTITION_TYPE_RB) -> list[list[str]]:
    nodes = get_nodes(name)
    return calculate_district_metering_area_for_nodes(name, nodes, part_count, part_type)