"""Provides the elements of a diagram."""

from typing import (
    Dict,
    Iterable,
    Iterator,
    List,
    Mapping,
    Optional,
    Sequence,
    Set,
    Tuple,
    Union,
)

from .attributes import (
    AttributeMap,
    Attributes,
    BlockAttributes,
    ConnectionAttributes,
    DiagramAttributes,
)

from .geometry import IntBounds, IntPoint, Orientation
from .util import log_info, log_warning

######################################################################

class RowDef:
    """Definition of a diagram row.

    Holds the information necessary to create a row of diagram cells.

    """

    def __init__(self, tags: Sequence[Optional[str]]):
        """Initialize with the given tags."""
        self._tags = tags

    @property
    def tags(self) -> Sequence[Optional[str]]:
        """Return the tags for the cells."""
        return list(self._tags)

######################################################################

class BlockDef:
    """Definition of a block.

    Holds the information necessary to create a block.

    """

    def __init__(
            self,
            name: Optional[str] = None,
            tags: Iterable[str] = (),
            **attrs: AttributeMap,
    ):
        """Initialize for a block with the given properties."""
        self._name = name
        tag_set: Set[str] = set()
        if tags:
            tag_set.update(tags)
        self._tags = tag_set
        self._attributes = attrs

    @property
    def name(self) -> Optional[str]:
        """Name of the block."""
        return self._name

    @property
    def tags(self) -> Set[str]:
        """Tags used to look up cells."""
        return set(self._tags)

    @property
    def attributes(self) -> AttributeMap:
        """The attributes to create the block with."""
        return self._attributes

######################################################################

# Connections can be made between whole blocks or nodes of blocks with
# a specifed tag.  A "sub-block" can be defined either by a block name
# (whole block) or a block name and a tag (partial block).
SubBlockDef = Union[str, Tuple[str, Optional[str]]]

######################################################################

class ConnectionDef:
    """Definition of a connection between two blocks.

    Holds the information necessary to create a connection.

    """

    def __init__(
            self,
            start: SubBlockDef,
            end: SubBlockDef,
            group: Optional[str] = None,
            **attrs: AttributeMap
    ):
        """Initialize given the names of two blocks."""
        self._start = start
        self._end = end
        self._group = group
        self._attributes = attrs

    @property
    def start(self) -> SubBlockDef:
        """Definition of source node."""
        return self._start

    @property
    def end(self) -> SubBlockDef:
        """Definition of destination node."""
        return self._end

    @property
    def group(self) -> Optional[str]:
        """Group to which the connection belongs."""
        return self._group

    @property
    def attributes(self) -> AttributeMap:
        """Attributes to create the connection with."""
        return self._attributes

######################################################################

# Multiple connection point can be specified as:
# - a single block name (string)
# - a list of block names (sequence of strings)
# - a list of block names and associated tag cells (mapping from
#   string to string)
MultipleNodes = Union[str, Sequence[str], Mapping[str, str]]

######################################################################

class DiagramDef:
    """Definition of a diagram.

    Holds the information necessary to create a diagram.

    """

    def __init__(self, **attrs: AttributeMap):
        """Initialize with the given diagram attributes.

        See DiagramAttributes for a list of available attributes.

        """
        self._attributes = Attributes(**attrs)
        self._auto_block_attributes = Attributes()
        self._row_defs: List[RowDef] = []
        self._block_defs: List[BlockDef] = []
        self._block_defs_by_name: Dict[str, BlockDef] = {}
        self._connection_defs: List[ConnectionDef] = []

    @property
    def attributes(self) -> Attributes:
        """The attributes to create the diagram with."""
        return self._attributes

    @property
    def auto_block_attributes(self) -> Attributes:
        """Attributes for blocks generated automatically from tags."""
        return self._auto_block_attributes

    def set_auto_block_attributes(self, **attrs: AttributeMap) -> None:
        """Set the attributes of the autogenerated blocks."""
        self._auto_block_attributes = Attributes(**attrs)

    def add_row(self, tags: Sequence[Optional[str]]) -> None:
        """Add a row at the end of the diagram.

        The input is a sequence of strings used to tag the cells with.
        An empty string or None results in an untagged cell.  The
        length of the new row is equal to the number of tags (empty or
        not).

        """
        rdef = RowDef(tags)
        self._row_defs.append(rdef)

    def add_block(
            self,
            name: Optional[str],
            tags: Iterable[str] = (),
            **attrs: AttributeMap,
    ) -> None:
        """Add a new block to the diagram.

        The 'tags' argument is a set of cell tags.  The block will
        cover the cells matching the tags in addition to the cells
        tagged with the name of the block itself.

        See BlockAttributes for a list of available attributes.

        Rejects the block with a warning if there is already a block
        registered with the same name.

        """
        by_name = self._block_defs_by_name
        if name and name in by_name:
            log_warning("Block '{}' already exists".format(name))
            return
        bdef = BlockDef(name, tags, **attrs)
        self._block_defs.append(bdef)
        if name:
            by_name[name] = bdef

    def add_connections(
            self,
            starts: MultipleNodes,
            ends: MultipleNodes,
            group: Optional[str] = None,
            **attrs: AttributeMap,
    ) -> None:
        """Create many connections at once.

        If the number of start blocks is n and the number of end
        blocks in m, then the number of connections created will be
        n*m (assuming all blocks exist).  See add_connection() for
        further information.

        """
        for start in self._to_sub_block_defs(starts):
            for end in self._to_sub_block_defs(ends):
                self.add_connection(start, end, group, **attrs)

    @staticmethod
    def _to_sub_block_defs(m: MultipleNodes) -> Sequence[SubBlockDef]:
        """Produce the object necessary for a connection."""
        if isinstance(m, str):
            return [m]
        elif isinstance(m, Sequence):
            return m
        elif isinstance(m, Mapping):
            seq = []
            for b, t in m.items():
                seq.append((b, t))
            return seq
        else:
            assert False, "Bad connection point definition"

    def add_connection(
            self,
            start: SubBlockDef, end: SubBlockDef,
            group: Optional[str] = None,
            **attrs: AttributeMap
    ) -> None:
        """Create a connection between two blocks.

        You must supply the start and the end of the connection either
        as block names or pairs of block name and cell tag.

        Provide the name of a group if you want the connection to
        belong to a group.

        See ConnectionAttributes for a list of available attributes.

        """
        cdef = ConnectionDef(start, end, group, **attrs)
        self._connection_defs.append(cdef)

    def row_defs(self) -> Iterator[RowDef]:
        """Return an iterator over the row definitions."""
        yield from self._row_defs

    def block_defs(self) -> Iterator[BlockDef]:
        """Return an iterator over the block definitions."""
        yield from self._block_defs

    def connection_defs(self) -> Iterator[ConnectionDef]:
        """Return an iterator over the connection definitions."""
        yield from self._connection_defs

######################################################################

class DiagramCell:
    """Position in the diagram grid."""

    def __init__(self, point: IntPoint, tag: Optional[str] = None):
        """Initialize for the given point with an optional tag."""
        self._point = point
        # Empty string is not a tag.
        if not tag:
            tag = None
        self._tag = tag

    def __repr__(self) -> str:
        """Convert to string."""
        return "{}(i={},j={},tag={})".format(
            self.__class__.__name__,
            self._point.i,
            self._point.j,
            self._tag,
        )

    @property
    def point(self) -> IntPoint:
        """Position of the object in the grid."""
        return self._point

    @property
    def tag(self) -> Optional[str]:
        """Tag used to look up the cell."""
        return self._tag

######################################################################

class DiagramRow:
    """A row in the diagram."""

    def __init__(self, i: int, tags: Sequence[Optional[str]]):
        """Initialize the row."""
        self._i = i
        cells = []
        for j, tag in enumerate(tags):
            p = IntPoint(i=i, j=j)
            cell = DiagramCell(p, tag)
            cells.append(cell)
        self._cells = cells

    def __repr__(self) -> str:
        """Convert to string."""
        return "{}(i={},width={})".format(
            self.__class__.__name__,
            self._i,
            len(self._cells),
        )

    def __len__(self) -> int:
        """Return the number of cells."""
        return len(self._cells)

    def __iter__(self) -> Iterator[DiagramCell]:
        """Return an iterator over the cells."""
        yield from self._cells

    def expand(self, n: int) -> None:
        """Add cells at the end if the row is shorter than the given length."""
        cells = self._cells
        m = len(cells)
        i = self._i
        for j in range(m, n):
            p = IntPoint(i=i, j=j)
            cell = DiagramCell(p)
            cells.append(cell)

######################################################################

class DiagramGrid:
    """Contains the cells of the diagram."""

    def __init__(self) -> None:
        """Initialize an empty grid."""
        self._rows: List[DiagramRow] = []

    def __iter__(self) -> Iterator[DiagramRow]:
        """Return an iterator over the rows."""
        yield from self._rows

    @property
    def height(self) -> int:
        """Number of rows."""
        return len(self._rows)

    @property
    def width(self) -> int:
        """Number of columns."""
        rows = self._rows
        if not rows:
            return 0
        else:
            return len(rows[0])

    def add_row(self, tags: Sequence[Optional[str]]) -> None:
        """Add a row of cells to the grid."""
        i = len(self._rows)
        row = DiagramRow(i, tags)
        rows = self._rows
        rows.append(row)
        # Adjust the width of the rows.
        new_width = max(self.width, len(row))
        for row in rows:
            row.expand(new_width)

    def cells_containing(self, tags: Iterable[str]) -> Iterator[DiagramCell]:
        """Return an iterator over the cells that contain the given tags.

        It calculates the minimum rectangular area that contains all
        the cells with the given tags.  The iterator iterates over all
        the cells in the rectangular area (i.e. not only the cells
        with the given tags).  It yields the cells in the order that
        they are stored in the grid, i.e. the order of the tags is
        irrelevant.

        """
        points = []
        for cell in self._cells_tagged(tags):
            points.append(cell.point)
        bounds = IntBounds.containing(points)
        if bounds:
            yield from self._cells_covering(bounds)

    def _cells_tagged(self, tags: Iterable[str]) -> Iterator[DiagramCell]:
        """Return an iterator over the cells with the given tags.

        The iterator yields the cells in the order that they are
        stored in the grid.  The order of the tags is irrelevant.

        """
        tag_set = set(tags)
        for cell in self._cells():
            tag = cell.tag
            if tag and tag in tag_set:
                yield cell

    def _cells_covering(self, bounds: IntBounds) -> Iterator[DiagramCell]:
        """Return an iterator over the cells covering the bounding box.

        The iterator yields the cells in the order that they are
        stored in the grid.

        """
        for row in self._rows:
            for cell in row:
                p = cell.point
                if (p.i >= bounds.imin and p.i <= bounds.imax and
                    p.j >= bounds.jmin and p.j <= bounds.jmax):
                    yield cell

    def _cells(self) -> Iterator[DiagramCell]:
        """Return an iterator over the cells.

        The iterator yields the cells in the order that they are
        stored in the grid.

        """
        for row in self._rows:
            yield from row

    def tags(self) -> Iterable[str]:
        """Return all the tags from the cells.

        This method ensures that the result contains each tag only
        once.  It returns the tags in no particular order.

        """
        tags = set()
        for cell in self._cells():
            tag = cell.tag
            if tag:
                tags.add(tag)
        return tags

######################################################################

class Node:
    """Endpoint of a connection on a block."""

    def __init__(self, cell: DiagramCell):
        """Initialize and place in the given cell."""
        self._cell = cell

    def __repr__(self) -> str:
        """Convert to string."""
        return "{}({})".format(self.__class__.__name__, self._cell)

    @property
    def point(self) -> IntPoint:
        """Position of the node in the diagram grid."""
        return self._cell.point

    @property
    def tag(self) -> Optional[str]:
        """Tag of the cell at the point."""
        return self._cell.tag

######################################################################

class Block:
    """Represents a block of the diagram."""

    def __init__(self, name: Optional[str] = None, **attrs: AttributeMap):
        """Initialize the block (with optional attributes)."""
        self._name = name
        self._attributes = BlockAttributes(**attrs)
        self._nodes: List[Node] = []
        self._bounds: Optional[IntBounds] = None

    def __repr__(self) -> str:
        """Convert to string."""
        return "{}({})".format(
            self.__class__.__name__,
            self._name,
        )

    @property
    def name(self) -> Optional[str]:
        """A name that identifies the block."""
        return self._name

    @property
    def attributes(self) -> BlockAttributes:
        """Attributes attached to the block."""
        return self._attributes

    @property
    def bounds(self) -> Optional[IntBounds]:
        """Return the bounding box of the block in the diagram grid."""
        b = self._bounds
        if b:
            return b.copy()
        else:
            return None

    def label(self) -> Optional[str]:
        """Return a label to draw on the block in the diagram."""
        label = self._attributes.label
        # If the label is not defined, use the name of the block.  If
        # we want a named block with an empty label, we have to use an
        # empty string for the label.  Thus, the empty string is valid
        # as a label.  Be specific here and compare with None.
        if label is None:
            return self._name
        else:
            return label

    def add_node(self, node: Node) -> None:
        """Associate the block with the node.

        It preserves the order of the nodes and ensures that no node
        is added more than once.

        """
        nodes = self._nodes
        if node not in nodes:
            nodes.append(node)
        # Recalculate the bounding box.
        self._update_bounds()

    def _update_bounds(self) -> None:
        """Update the bounding box covered by the nodes."""
        points = []
        for node in self._nodes:
            points.append(node.point)
        self._bounds = IntBounds.containing(points)

    def nodes(self) -> Iterator[Node]:
        """Return an iterator over the nodes."""
        yield from self._nodes

    def overlaps_with(self, other: 'Block') -> bool:
        """True if the two blocks have common nodes."""
        nodes_1 = set(self.nodes())
        nodes_2 = set(other.nodes())
        if nodes_1.intersection(nodes_2):
            return True
        else:
            return False

######################################################################

class SubBlock:
    """Part of a block that consists of specific nodes."""

    def __init__(self, block: Block, tag: Optional[str] = None):
        """Initialize for a given block and optional cell tag."""
        self._block = block
        self._tag = tag
        nodes: List[Node] = []
        for node in block.nodes():
            if not tag or node.tag == tag:
                nodes.append(node)
        self._nodes = nodes

    def __len__(self) -> int:
        """Return the number of nodes."""
        return len(self._nodes)

    def __repr__(self) -> str:
        """Convert to string."""
        return "{}({}:{})".format(
            self.__class__.__name__,
            self._block.name,
            self._tag,
        )

    @property
    def block(self) -> Block:
        """The block."""
        return self._block

    @property
    def name(self) -> Optional[str]:
        """Name of the block."""
        return self._block.name

    def outer_nodes(self) -> Iterable[Node]:
        """Return the subset of the nodes at the sides of the block."""
        bounds = self._block.bounds
        if bounds:
            for node in self._nodes:
                i, j = node.point
                if (i == bounds.imin or i == bounds.imax or
                    j == bounds.jmin or j == bounds.jmax):
                    yield node

######################################################################

class Connection:
    """A connection between two blocks in the diagram."""

    def __init__(
            self,
            start: SubBlock, end: SubBlock,
            group: Optional[str] = None,
            **attrs: AttributeMap
    ):
        """Initialize a connection between the given nodes."""
        self._start = start
        self._end = end
        self._group = group
        self._attributes = ConnectionAttributes(**attrs)

    def __repr__(self) -> str:
        """Convert to string."""
        return "{}({}->{})".format(
            self.__class__.__name__,
            self._start,
            self._end,
        )

    @property
    def start(self) -> SubBlock:
        """Source block of the connection."""
        return self._start

    @property
    def end(self) -> SubBlock:
        """Destination block of the connection."""
        return self._end

    @property
    def group(self) -> Optional[str]:
        """Group to which the connection belongs."""
        return self._group

    @property
    def attributes(self) -> ConnectionAttributes:
        """Attributes attached to the connection."""
        return self._attributes

######################################################################

class Diagram:
    """Container for the blocks and connections of the diagram."""

    def __init__(self, ddef: DiagramDef):
        """Initialize from the definition."""
        self._attributes = DiagramAttributes(**ddef.attributes)
        self._grid = self._make_grid(ddef)
        self._blocks: List[Block] = []
        self._blocks_by_name: Dict[str, Block] = {}
        self._nodes_to_blocks: Dict[Node, List[Block]] = {}
        self._points_to_nodes: Dict[IntPoint, Node] = {}
        self._init_blocks(ddef)
        self._connections: List[Connection] = []
        self._init_connections(ddef)

    @staticmethod
    def _make_grid(ddef: DiagramDef) -> DiagramGrid:
        """Create the grid from the diagram definition."""
        grid = DiagramGrid()
        for row_def in ddef.row_defs():
            grid.add_row(row_def.tags)
        return grid

    def _init_blocks(self, ddef: DiagramDef) -> None:
        """Create the blocks from the definition."""
        blocks = self._blocks
        blocks.clear()
        nodes_to_blocks = self._nodes_to_blocks
        nodes_to_blocks.clear()
        points_to_nodes = self._points_to_nodes
        points_to_nodes.clear()
        self._make_defined_blocks(ddef)
        self._make_leftover_blocks(ddef)

    def _make_defined_blocks(self, ddef: DiagramDef) -> None:
        """Create the blocks defined in the diagram definition."""
        self._make_blocks(ddef.block_defs())

    def _make_leftover_blocks(self, ddef: DiagramDef) -> None:
        """Create blocks from leftover tags."""
        attrs = ddef.auto_block_attributes
        tags = sorted(self._leftover_tags(ddef))
        bdefs = []
        for tag in tags:
            bdef = BlockDef(tag, [], **attrs)
            bdefs.append(bdef)
        self._make_blocks(bdefs)

    def _leftover_tags(self, ddef: DiagramDef) -> Set[str]:
        """Cell tags that do not appear in block definitions."""
        tags = set(self._grid.tags())
        for bdef in ddef.block_defs():
            block_tags = set(bdef.tags)
            name = bdef.name
            if name:
                block_tags.add(name)
            for tag in block_tags:
                if tag in tags:
                    tags.remove(tag)
        return tags

    def _make_blocks(self, bdefs: Iterable[BlockDef]) -> None:
        """Use the block definitions to create blocks."""
        blocks = self._blocks
        by_name = self._blocks_by_name
        nodes_to_blocks = self._nodes_to_blocks
        points_to_nodes = self._points_to_nodes
        grid = self._grid
        for bdef in bdefs:
            name = bdef.name
            block = Block(name, **bdef.attributes)
            blocks.append(block)
            tags = set()
            if name:
                by_name[name] = block
                # The name of the block is a tag itself!
                tags.add(name)
            tags.update(bdef.tags)
            # Find or create the nodes.
            for cell in grid.cells_containing(tags):
                p = cell.point
                node = points_to_nodes.get(p)
                if not node:
                    node = Node(cell)
                    points_to_nodes[p] = node
                block.add_node(node)
                if node not in nodes_to_blocks:
                    nodes_to_blocks[node] = []
                nodes_to_blocks[node].append(block)

    def _init_connections(self, ddef: DiagramDef) -> None:
        """Create the connections from the definition."""
        connections = self._connections
        connections.clear()
        for cdef in ddef.connection_defs():
            connection = self._make_connection(cdef)
            if connection:
                connections.append(connection)

    def _make_connection(self, cdef: ConnectionDef) -> Optional[Connection]:
        """Create a connection from a connection definition."""
        start = self._sub_block(cdef.start)
        end = self._sub_block(cdef.end)
        if not (start and end):
            return None
        # Cannot create connection between overlapping blocks.
        block1 = start.block
        block2 = end.block
        if block1.overlaps_with(block2):
            tmpl = "Blocks '{}' and '{}' overlap, connection rejected"
            log_warning(tmpl.format(block1.name, block2.name))
            return None
        # Everything seems OK, let's make the connection.
        connection = Connection(start, end, cdef.group, **cdef.attributes)
        return connection

    def _sub_block(self, ndef: SubBlockDef) -> Optional[SubBlock]:
        """Retrieve the subset of the block from the definition."""
        if isinstance(ndef, str):
            block_name = ndef
            tag = None
        else:
            block_name, tag = ndef
        block = self._block(block_name)
        if block:
            sub = SubBlock(block, tag)
            if sub:
                return sub
            else:
                msg = "Node '{}' is not placed""".format(block_name)
                log_warning(msg)
                return None
        else:
            return None

    def _block(self, name: str) -> Optional[Block]:
        """Retrieve a block by name."""
        block = self._blocks_by_name.get(name)
        if not block:
            log_warning("Block '{}' does not exist".format(name))
        return block

    @property
    def grid(self) -> DiagramGrid:
        """The grid that contains the blocks."""
        return self._grid

    @property
    def attributes(self) -> DiagramAttributes:
        """Attributes attached to the diagram."""
        return self._attributes

    def blocks(self) -> Iterator[Block]:
        """Return an iterator over the blocks.

        The iterator yields the blocks in the order they were added to
        the diagram.

        """
        yield from self._blocks

    def node_blocks(self, node: Node) -> Iterator[Block]:
        """Iterate over the blocks connected to a node."""
        yield from self._nodes_to_blocks[node]

    def nodes_and_blocks(self) -> Iterator[Tuple[Node, Set[Block]]]:
        """Return an iterator over the nodes and the blocks on it."""
        for node, blocks in self._nodes_to_blocks.items():
            yield node, set(blocks)

    def connections(self) -> Iterator[Connection]:
        """Return an iterator over the connections."""
        yield from self._connections

    def _pretty_print(self) -> None:
        """Print the diagram for debugging purposes."""
        print("Blocks:")
        for block in self._blocks:
            print("\t{}".format(block))
        print("Connections:")
        for connection in self._connections:
            print("\t{}".format(connection))
        print("Rows:")
        for row in self._grid:
            print("\t{}:".format(row))
            for cell in row:
                print("\t\t{}".format(cell))
