"""Lay out expression dependencies with ordered operands and reserved wire passages."""

from collections import defaultdict
from dataclasses import dataclass
from graphlib import TopologicalSorter

import numpy as np
from scipy.optimize import linprog

from .flow import peer_links


@dataclass
class FieldLayout:
    positions: dict
    width: float
    height: float
    routes: list


def field_layout(nodes, parents, peers, sizes, *, gap, target_height):
    """Order dependency layers, reserve long-link passages, and solve socket heights."""
    roots = {node: node for node in nodes}

    def root(node):
        while roots[node] != node:
            node = roots[node]
        return node

    for sources in peers:
        members = [node for node in sources if node in roots]
        if sum(sizes[node].height for node in members) + gap * (len(members) - 1) > target_height:
            continue
        for node in members[1:]:
            roots[root(node)] = root(members[0])
    groups = defaultdict(list)
    for node in nodes:
        groups[root(node)].append(node)
    for sources in peers:
        members = [node for node in sources if node in roots]
        if members and all(root(node) == root(members[0]) for node in members):
            groups[root(members[0])].sort(key=lambda node: sources.index(node) if node in sources else len(sources))
    blocks = list(groups.values())
    block_of = {node: index for index, block in enumerate(blocks) for node in block}
    children = {index: set() for index in range(len(blocks))}
    for node in nodes:
        for parent in parents[node]:
            if parent in block_of and block_of[parent] != block_of[node]:
                children[block_of[parent]].add(block_of[node])
    ranks = {}
    for block in TopologicalSorter(children).static_order():
        ranks[block] = max((ranks[child] + 1 for child in children[block]), default=0)
    depth = max(ranks.values())
    ranks = {block: depth - rank for block, rank in ranks.items()}
    heights = [sum(sizes[node].height for node in block) + gap * (len(block) - 1) for block in blocks]
    offsets = {}
    for index, block in enumerate(blocks):
        cursor = -heights[index] / 2
        for node in block:
            offsets[node] = cursor
            cursor += sizes[node].height + gap
    layers = defaultdict(list)
    for block, rank in ranks.items():
        layers[rank].append(block)
    operand_columns = {rank for node in nodes for link in peer_links(node)
                       if link.from_node in block_of
                       for rank in range(ranks[block_of[link.from_node]], ranks[block_of[node]] + 1)}
    # Keep an operand block and its junction together when spreading tall layers.
    columns, column = defaultdict(list), 0
    for rank in sorted(layers):
        used = 0
        for block in layers[rank]:
            if rank not in operand_columns and used and used + gap + heights[block] > target_height:
                column += 1
                used = 0
            columns[column].append(block)
            ranks[block] = column
            used += heights[block] + (gap if used else 0)
        column += 1
    layers = columns
    depth = column - 1
    links = [link for node in nodes for socket in node.inputs for link in socket.links if link.from_node in block_of]
    edges, passages = [], []
    for link in links:
        first, last = block_of[link.from_node], block_of[link.to_node]
        if first == last:
            raise ValueError("Equal operands must not contain an internal dependency")
        start_offset = offsets[link.from_node] + sizes[link.from_node].sockets[link.from_socket]
        end_offset = offsets[link.to_node] + sizes[link.to_node].sockets[link.to_socket]
        previous, previous_offset, path = first, start_offset, []
        for rank in range(ranks[first] + 1, ranks[last]):
            point = len(heights)
            heights.append(16)
            ranks[point] = rank
            layers[rank].append(point)
            edges.append((previous, point, previous_offset, 0))
            previous, previous_offset = point, 0
            path.append(point)
        edges.append((previous, last, previous_offset, end_offset))
        if path:
            passages.append((link, path))
    incoming, outgoing = defaultdict(list), defaultdict(list)
    for first, last, a, b in edges:
        outgoing[first].append((last, a, b))
        incoming[last].append((first, b, a))
    centers = {}

    def stack(layer):
        height = sum(heights[item] for item in layer) + sum(separation(a, b) for a, b in zip(layer, layer[1:]))
        cursor = -height / 2
        for index, item in enumerate(layer):
            centers[item] = cursor + heights[item] / 2
            cursor += heights[item] + (separation(item, layer[index + 1]) if index + 1 < len(layer) else 0)

    def separation(first, last):
        return gap if first < len(blocks) and last < len(blocks) else gap / 2

    for layer in layers.values():
        stack(layer)
    for _ in range(4):
        for ranks_in_order, adjacent in ((range(depth + 1), incoming), (range(depth, -1, -1), outgoing)):
            for rank in ranks_in_order:
                layer = layers[rank]
                def desired(item):
                    neighbors = adjacent[item]
                    return sum(centers[other] + their_offset - own_offset
                               for other, own_offset, their_offset in neighbors) / len(neighbors) if neighbors else centers[item]
                layer.sort(key=lambda item: (desired(item), item))
                stack(layer)

    count = len(heights)
    variable_count = count + 1 + len(edges)
    inequalities, bounds, equalities, values = [], [], [], []

    def constraint(terms, maximum):
        row = np.zeros(variable_count)
        for index, coefficient in terms:
            row[index] += coefficient
        inequalities.append(row)
        bounds.append(maximum)

    for index, height in enumerate(heights):
        constraint([(index, -1)], -height / 2)
        constraint([(index, 1), (count, -1)], -height / 2)
    for layer in layers.values():
        for first, last in zip(layer, layer[1:]):
            constraint([(first, 1), (last, -1)], -(heights[first] + heights[last]) / 2 - separation(first, last))
    for index, (first, last, a, b) in enumerate(edges):
        slack = count + 1 + index
        constraint([(first, 1), (last, -1), (slack, -1)], b - a)
        constraint([(last, 1), (first, -1), (slack, -1)], a - b)
    for node in nodes:
        inputs = [link.from_socket for link in peer_links(node)]
        if len(inputs) < 2 or any(socket.node not in block_of for socket in inputs):
            continue
        output = next(socket for socket in node.outputs if socket in sizes[node].sockets)
        first, last = inputs[0], inputs[-1]
        row = np.zeros(variable_count)
        row[block_of[node]] += 2
        row[block_of[first.node]] -= 1
        row[block_of[last.node]] -= 1
        equalities.append(row)
        values.append(offsets[first.node] + sizes[first.node].sockets[first]
                      + offsets[last.node] + sizes[last.node].sockets[last]
                      - 2 * (offsets[node] + sizes[node].sockets[output]))
    objective = np.zeros(variable_count)
    objective[count] = len(edges) + 1
    objective[count + 1:] = 1
    result = linprog(objective, A_ub=np.array(inequalities), b_ub=bounds,
                     A_eq=np.array(equalities) if equalities else None, b_eq=values or None,
                     bounds=[(None, None)] * count + [(0, None)] * (1 + len(edges)), method="highs")
    if not result.success:
        raise ValueError(f"Expression layers could not be placed: {result.message}")
    centers = result.x[:count]
    positions, point_positions, right = {}, {}, 0
    for rank in range(depth + 1):
        layer = layers[rank]
        width = max(max(sizes[node].width for node in blocks[item]) if item < len(blocks) else 16 for item in layer)
        if rank:
            run = max([gap, *(abs(centers[first] + a - centers[last] - b) * 0.8
                             for first, last, a, b in edges if ranks[last] == rank)])
            right += run
        right += width
        for item in layer:
            if item < len(blocks):
                for node in blocks[item]:
                    positions[node] = (right - sizes[node].width, -centers[item] - offsets[node])
            else:
                point_positions[item] = (right, -centers[item])
    routes = []
    for link, path in passages:
        points = [point_positions[item] for item in path]
        points = [point for index, point in enumerate(points) if index == 0 or index == len(points)-1
                  or not abs(points[index-1][1] - point[1]) < 1e-6 or not abs(points[index+1][1] - point[1]) < 1e-6]
        routes.append((link, points))
    return FieldLayout(positions, right, result.x[count], routes)