from .database import ChangeSet
from .s0_base import is_junction, get_nodes
from .s9_demands import get_demand
from .s32_region_util import Topology, get_nodes_in_region
from .batch_exe import execute_batch_command


DISTRIBUTION_TYPE_ADD = 'ADD'
DISTRIBUTION_TYPE_OVERRIDE = 'OVERRIDE'


def calculate_demand_to_nodes(name: str, demand: float, nodes: list[str]) -> dict[str, float]:
    if len(nodes) == 0 or demand == 0.0:
        return {}

    topology = Topology(name, nodes)
    t_nodes = topology.nodes()
    t_links = topology.links()

    length_sum = 0.0
    for value in t_links.values():
        length_sum += abs(value['length'])

    if length_sum <= 0.0:
        return {}

    demand_per_length = demand / length_sum

    result: dict[str, float] = {}
    for node, value in t_nodes.items():
        if not is_junction(name, node):
            continue
        demand_per_node = 0.0
        for link in value['links']:
            demand_per_node += abs(t_links[link]['length']) * demand_per_length * 0.5
        result[node] = demand_per_node

    return result


def calculate_demand_to_region(name: str, demand: float, region: str) -> dict[str, float]:
    nodes = get_nodes_in_region(name, region)
    return calculate_demand_to_nodes(name, demand, nodes)


def calculate_demand_to_network(name: str, demand: float) -> dict[str, float]:
    nodes = get_nodes(name)
    return calculate_demand_to_nodes(name, demand, nodes)


def distribute_demand_to_nodes(name: str, demand: float, nodes: list[str], type: str = DISTRIBUTION_TYPE_ADD) -> ChangeSet:
    if len(nodes) == 0 or demand == 0.0:
        return ChangeSet()
    if type != DISTRIBUTION_TYPE_ADD and type != DISTRIBUTION_TYPE_OVERRIDE:
        return ChangeSet()

    topology = Topology(name, nodes)
    t_nodes = topology.nodes()
    t_links = topology.links()

    length_sum = 0.0
    for value in t_links.values():
        length_sum += abs(value['length'])
    
    if length_sum <= 0.0:
        return ChangeSet()
    
    demand_per_length = demand / length_sum

    cs = ChangeSet()

    for node, value in t_nodes.items():
        if not is_junction(name, node):
            continue
        demand_per_node = 0.0
        for link in value['links']:
            demand_per_node += abs(t_links[link]['length']) * demand_per_length * 0.5

        ds = get_demand(name, node)['demands']
        if len(ds) == 0:
            ds = [{'demand': demand_per_node, 'pattern': None, 'category': None}]
        elif type == DISTRIBUTION_TYPE_ADD:
            ds[0]['demand'] += demand_per_node
        else:
            ds[0]['demand'] = demand_per_node
        cs.update({'type': 'demand', 'junction': node, 'demands': ds})

    return execute_batch_command(name, cs)


def distribute_demand_to_region(name: str, demand: float, region: str, type: str = DISTRIBUTION_TYPE_ADD) -> ChangeSet:
    nodes = get_nodes_in_region(name, region)
    return distribute_demand_to_nodes(name, demand, nodes, type)