import os

import matplotlib.pyplot as plt
import networkx as nx
from IPython.core.display import SVG
from pyflowsheet import (
    BlackBox,
    Distillation,
    Flowsheet,
    HeatExchanger,
    HorizontalLabelAlignment,
    StreamFlag,
    SvgContext,
    VerticalLabelAlignment,
    Vessel,
)
from tabulate import tabulate


def create_stream_table(graph, chemicalspecies, decimals):
    """Function to return a table with all information about the streams in the flowsheet.

    Args:
        graph: nx-Graph representation of the flowsheet
        chemicalspecies: List with string names of chemical species considered in the process
        decimals: number of decimals showed int the table

    Returns:
        table_streams: table with all streams
    """
    if chemicalspecies is None:
        chemicalspecies = ["x_A", "x_B", "x_C"]

    header = ["Stream name", "Edge", "Moles(mol/s)", "Temperature(Kelvin)", "Pressure(Pa)"]
    header.extend(chemicalspecies)

    table_data = [header]
    for edge in graph.edges:
        stream_list = []
        stream_list.append(graph.get_edge_data(*edge)["processstream_name"])
        stream_data = graph.get_edge_data(*edge)["processstream_data"]
        stream_list.append(edge)
        for i in range(len(stream_data[0:3])):
            stream_list.append(round(stream_data[i], decimals))
        for i in range(len(stream_data[-1])):
            stream_list.append(round(stream_data[-1][i], decimals))
        table_data.append(stream_list)
    table_streams = tabulate(table_data, headers="firstrow", tablefmt="grid")
    print(table_streams)
    return table_streams


def create_unit_table(graph, decimals):
    """Function to return a table with all information about the units in the flowsheet.

    Args:
        graph: nx-Graph representation of the flowsheet
        decimals: number of decimals showed int the table

    Returns:
        table_units: table with all units
    """
    unit_table_data = [
        ["Unit name", "Unit type", "Condition description", "Condition"]]
    for node in graph.nodes:
        node_list = []
        node_list.append(node)
        node_list.append(graph.nodes(data=True)[node]["unit_type_specific"])
        unit = graph.nodes(data=True)[node]["unit"]
        if graph.nodes(data=True)[node]["unit_type"] == "hex":
            node_list.append("water inlet temperature")
            node_list.append(round(unit.water_temp_in, decimals))
        elif graph.nodes(data=True)[node]["unit_type"] == "r":
            node_list.append("length")
            node_list.append(round(unit.length, decimals))
        elif graph.nodes(data=True)[node]["unit_type"] == "col":
            node_list.append("distillation to feed ratio")
            node_list.append(round(unit.has_distillation_to_feed_ratio, decimals))
        elif graph.nodes(data=True)[node]["unit_type"] == "splt":
            node_list.append("split ratio")
            node_list.append(round(unit.split_ratio, decimals))
        else:
            node_list.append("N/A")
            node_list.append(0)
        unit_table_data.append(node_list)
    table_units = tabulate(unit_table_data, headers="firstrow", tablefmt="grid")
    print(table_units)
    return table_units


def _add_positions(graph, flowsheet_size):
    """Function to assign positions to each node of a flowsheet graph.

    Args:
        graph: nx-graph representation of the flowsheet
        flowsheet_size: number of nodes in the flowsheet

    Returns:
        graph with updated node attributes
    """
    # First set position of feed(s)

    # Save all feed nodes in a list (reason: the following algorithm always checks for all outgoing edges
    # and their successor nodes, starting with the first feeds. To come back to those later, they and their
    # position are saved in lists
    save_nodes = [feed for feed in graph.nodes if graph.in_degree(feed) == 0]
    save_pos = []
    y_coordinates = []  # List to save all y coordinates --> no nodes above each other
    # Initialize a list for saving all nodes that have been updated
    updated_nodes = []
    i=0
    for _, feed in enumerate(save_nodes):
        # Update positions of all feeds
        nx.set_node_attributes(graph, {feed: {"pos": [0, i]}})
        y_coordinates.append(i)
        # Save position of all feeds
        save_pos.append([0, i])
        # Write feeds in updated_nodes list to check later, if they already have a position attribute
        # (necessary when dealing with recycles!)
        updated_nodes.append(feed)
        i += 500

    # For the following loop, start from the first feed
    node = save_nodes[0]
    pos = save_pos[0]
    # Initialize a counter (set to one because the first feed is already in use)
    counter = 1

    while len(updated_nodes) < flowsheet_size:
        # While loop for as long as not all units have been updated with a position
        # Find all edges that leave the previously set node
        original_edges = list(graph.out_edges(node))
        # FOR RECYCLES: only considere edges that lead to nodes that don't have a position yet!
        # Copy all relevant edges to the list edges
        edges = []
        for k in range(len(original_edges)):
            if original_edges[k][1] not in updated_nodes:
                edges.append(original_edges[k])

        # Set position of the following node(s) depending on their number
        if len(edges) == 1:
            # Only one next node
            next_node = edges[0][1]
            # Set position to the right
            pos = pos[:]
            pos[0] += 150
            nx.set_node_attributes(graph, {next_node: {"pos": pos}})
            # Write updated node into list to check later, if this node has already been updated
            updated_nodes.append(next_node)
            # Use the new node as the source-node for the next iteration
            node = next_node

        elif len(edges) == 2:
            # Two successor nodes: Positions are set to the right and up / down
            next_node = edges[0][1]
            pos = pos[:]
            pos_x = pos[0]
            pos_y = pos[1]
            pos[0] = pos_x + 200
            pos[1] = pos_y + 100
            multiplier = 0.5

            while pos[1] in y_coordinates:
                # Quick fix to prevent multiple nodes directly above each other --> for side-streams a new y-coordinate
                # is chosen that has not been used before
                pos[1] = pos_y + 100 * multiplier
                multiplier = multiplier * 0.5

            y_coordinates.append(pos[1])
            nx.set_node_attributes(graph, {next_node: {"pos": pos}})
            updated_nodes.append(next_node)
            # Save the first of the two nodes and its position in the save lists in order to come back to
            # them later
            save_nodes.append(next_node)
            save_pos.append(pos)

            # Update the second node just like in the case with only one successor node
            next_node = edges[1][1]
            pos = pos[:]
            pos[1] -= 200
            multiplier = 0.5

            while pos[1] in y_coordinates:
                pos[1] = pos_y - 150 * multiplier
                multiplier = multiplier * 0.5

            y_coordinates.append(pos[1])
            nx.set_node_attributes(graph, {next_node: {"pos": pos}})
            updated_nodes.append(next_node)
            node = next_node

        elif len(edges) > 2:
            # Multiple successor nodes: Positions are set to the right and down
            
            for i in range(len(edges)):
                next_node = edges[i][1]
                pos = pos[:]
                pos_x = pos[0]
                pos_y = pos[1]
                pos[0] = pos_x + 200
                pos[1] = pos_y - 200 * i
                multiplier = 0.5

                while pos[1] in y_coordinates:
                    pos[1] = pos_y - 150 * multiplier
                    multiplier = multiplier * 0.5

                y_coordinates.append(pos[1])
                nx.set_node_attributes(graph, {next_node: {"pos": pos}})
                updated_nodes.append(next_node)
                if i == len(edges) - 1:
                    node = next_node
                else:
                    # Save all but the last nodes and their position in the save lists in order to come back to
                    # them later
                    save_nodes.append(next_node)
                    save_pos.append(pos)

        elif len(edges) == 0:
            # If no edge is coming out of the node --> go on with the nodes that have been saved in the
            # save list and use them for the next iteration.
            node = save_nodes[counter]
            pos = save_pos[counter]
            counter += 1

    return graph


def plot_flowsheet_nx(graph, plot_with_stream_labels, add_positions=True):
    """Function to return a plot of a flowsheet represented as nx-graph. The visualization of the nodes is oriented on
    real flowsheets (start at the left side with the feed and move towards the products on the right).

    Args:
        graph: nx-Graph representation of the flowsheet
        plot_with_stream_labels: boolean telling whether stream data should be plotted or not
        add_positions: boolean telling whether a position attribute already exists

    Returns:
        table_streams: table with all streams
    """
    # Plots the flowsheet using matplotlib and shows the moles, temperature and pressure for each stream.
    # Update position attribute of nodes
    flowsheet_size = graph.number_of_nodes()
    fig = plt.figure(figsize=(flowsheet_size * 10. / 6., flowsheet_size * 5. / 6.))

    if add_positions:
        graph = _add_positions(graph, flowsheet_size)
        pos = nx.get_node_attributes(graph, "pos")
    else:
        pos = nx.get_node_attributes(graph, "pos")
        max_x = max([pos[node][0] for node in pos])
        max_y = max([pos[node][1] for node in pos])
        fig = plt.figure(figsize=(max_x/100, max_y/100))

    # Plot
    if plot_with_stream_labels:
        nx.draw(graph, pos, with_labels=True, node_size=1600, font_size=13, node_color="#00b4d9")
        try:
            labels = dict([((n1, n2),
                            "".join([d["processstream_name"], "\n N=", str(round(d["processstream_data"][0])),
                                     " mol/s\nT=", str(round(d["processstream_data"][1])), " K\nP=",
                                     str(round(d["processstream_data"][2])), " Pa"]))
                           for n1, n2, d in graph.edges(data=True)])
            nx.draw_networkx_edge_labels(graph, pos, edge_labels=labels, rotate=False, font_size=10)

        except KeyError:
            print("Key error! stream does not have all attributes")
        except IndexError:
            print("Index error! stream does not have all attributes")

    else:
        nx.draw(graph, pos, with_labels=True, node_size=1600, font_size=13, node_color="#00b4d9")
    plt.show()
    return fig


def plot_flowsheet_pyflowsheet(graph, block=False, imagepath="flowsheet",  pfd_id="PFD",
                               pfd_name="process flow diagram", pfd_description="created with pyflowsheet",
                               add_positions=True):
    """Function to plot a flowsheet-graph using the package pyflowsheet.

    Args:
        graph: nx-graph representation of the flowsheet
        block (boolean): True: plot a block flowsheet, False: plot a "normal" flowsheet with images as units
        imagepath (string): path where the svg is saved, e.g. plots/blockflowprocess
        pfd_id (string): ID of the flowsheet
        pfd_name (string): Name of the flowsheet
        pfd_description (string): Description of the flowsheet
        add_positions: boolean telling whether a position attribute already exists
    """
    # First assign a position to every node
    flowsheet_size = graph.number_of_nodes()

    if add_positions:
        graph = _add_positions(graph, flowsheet_size)

    # create a flowsheet
    pfd = Flowsheet(pfd_id, pfd_name, pfd_description)  # id, name, description

    # Intialize a dict to store all unit operations using the node-ids as keys
    unit_dict = dict.fromkeys(list(graph.nodes))

    # Loop over all nodes in the graph to create UnitOperation objects with the pyflowsheet package
    feed_count = 1
    product_count = 1
    for node_id, node in graph.nodes(data=True):
        # Find all feeds and assign them to a StreamFlag unit
        if graph.in_degree(node_id) == 0:
            feed_name = "Feed " + str(feed_count)
            feed = StreamFlag(node_id, name=feed_name, position=node["pos"])
            feed.setTextAnchor(HorizontalLabelAlignment.Center, VerticalLabelAlignment.Center, (0, 5))  # Text in image
            feed_count += 1
            unit_dict[node_id] = feed

        # Find all products and assign them to a StreamFlag unit
        elif graph.out_degree(node_id) == 0 and node_id[0] == "I":
            # Second condition in case a flowsheet does not end with IO unit but e.g. with 'X' (qick fix, there must be
            # a better way)
            product_name = "Product " + str(product_count)
            product = StreamFlag(node_id, name=product_name, position=node["pos"])
            product.setTextAnchor(HorizontalLabelAlignment.Center, VerticalLabelAlignment.Center, (0, 5))
            product_count += 1
            unit_dict[node_id] = product

        else:  # Unit operations
            if block:  # In a block-flowsheet all units are represented by a box
                unit = BlackBox(node_id, name=node_id, size=(80, 60), position=node["pos"])
                unit.setTextAnchor(HorizontalLabelAlignment.Center, VerticalLabelAlignment.Center, (0, 5))
            else:  # Use images for units
                # Todo: is there a way to work around the if-else statement?
                if node["unit_type"] == "hex":
                    unit = HeatExchanger(node_id, name=node_id, position=node["pos"])
                elif node["unit_type"] == "r":
                    unit = Vessel(node_id, name=node_id, position=node["pos"], angle=90)
                    unit.setTextAnchor(HorizontalLabelAlignment.Center, VerticalLabelAlignment.Center, (0, 5))
                elif node["unit_type"] == "col":
                    unit = Distillation(node_id, name=node_id, position=node["pos"], hasReboiler=False,
                                        hasCondenser=False)
                    unit.setTextAnchor(HorizontalLabelAlignment.Center, VerticalLabelAlignment.Center, (0, 5))
                else:
                    unit = BlackBox(node_id, name=node_id, position=node["pos"], size=(80, 60))
                    unit.setTextAnchor(HorizontalLabelAlignment.Center, VerticalLabelAlignment.Center, (0, 5))
            unit_dict[node_id] = unit

    # Connect all units
    if block:  # All ports are called 'In' or 'Out' --> makes it easier
        count = 1  # Count all streams for identifier
        for edge in graph.out_edges(data=True):
            unit_1 = unit_dict[edge[0]]
            unit_2 = unit_dict[edge[1]]
            stream_id = "stream-"+str(count)
            pfd.connect(stream_id, unit_1["Out"], unit_2["In"])

            # Set stream name position
            pos0 = graph.nodes(data=True)[edge[0]]["pos"]
            pos1 = graph.nodes(data=True)[edge[1]]["pos"]
            if pos0[1] > pos1[1]:
                pfd.streams[stream_id].labelOffset = (15, 10)
            else:
                pfd.streams[stream_id].labelOffset = (15, -10)
            count += 1

    else:  # If other UnitOperation objects than "BlackBox" are used, all ports have different names
        count = 1
        for edge in graph.out_edges(data=True):
            unit_1 = unit_dict[edge[0]]
            unit_2 = unit_dict[edge[1]]
            stream_id = "stream-" + str(count)
            pos0 = graph.nodes(data=True)[edge[0]]["pos"]
            pos1 = graph.nodes(data=True)[edge[1]]["pos"]

            # Save names of ports of each specific unit
            # Todo: is there a way to work around the if-else statement?
            if graph.nodes(data=True)[edge[0]]["unit_type"] == "hex":
                port1 = "TOut"
            elif graph.nodes(data=True)[edge[0]]["unit_type"] == "col":
                if pos0[1] > pos1[1]:
                    port1 = "LOut"
                else:
                    port1 = "VOut"
            else:
                port1 = "Out"
            if graph.nodes(data=True)[edge[1]]["unit_type"] == "hex":
                port2 = "TIn"
            elif graph.nodes(data=True)[edge[1]]["unit_type"] == "col":
                port2 = "Feed"
            else:
                port2 = "In"

            pfd.connect(stream_id, unit_1[port1], unit_2[port2])

            # Set stream name position
            if pos0[1] > pos1[1]:
                pfd.streams[stream_id].labelOffset = (25, 20)
            else:
                pfd.streams[stream_id].labelOffset = (25, -20)
            count += 1

    # Add units to pfd
    pfd.addUnits(unit_dict.values())

    # store the flowsheet as svg
    filename = imagepath + ".svg"
    ctx = SvgContext(filename)
    img = pfd.draw(ctx)

    if not os.path.exists("plots"):
        os.makedirs("plots")

    SVG(img.render(scale=1))
