#!/usr/bin/env python3
"""
This script provides a temporary solution for the problem of inout top level
port representation of the BLIF format.

The script loads the design from a JSON file generated by Yosys. Then it splits
all inout ports along with their nets and connections to cell ports into two.
Suffixes are automatically added to distinguish between the input and the
output part.

The JSON design format used by Yosys is documented there:
- http://www.clifford.at/yosys/cmd_write_json.html
- http://www.clifford.at/yosys/cmd_read_json.html

For example in the given design (in verilog):

 module top(
  input  A,
  output B,
  inout  C
 );

  IOBUF iobuf (
   .I(A),
   .O(B),
   .IO_$inp(C),
   .IO_$out(C)
  );

 endmodule

the resulting design would be:

 module top(
  input  A,
  output B,
  input  C_$inp,
  output C_$out
 );

  IOBUF iobuf (
   .I(A),
   .O(B),
   .IO_$inp(C_$inp),
   .IO_$out(C_$out)
  );

 endmodule

"""
import argparse
import os
import simplejson as json

# =============================================================================


def find_top_module(design):
    """
    Looks for the top-level module in the design. Returns its name. Throws
    an exception if none was found.
    """

    for name, module in design["modules"].items():
        attrs = module["attributes"]
        if "top" in attrs and int(attrs["top"]) == 1:
            return name

    raise RuntimeError("No top-level module found in the design!")


def get_nets(bits):
    """
    Returns a set of numbers corresponding to net indices effectively skipping
    connections to consts ("0", "1", "x").

    >>> get_nets([0, 1, 2, "0", "1", "x", 3, 4, 5])
    {0, 1, 2, 3, 4, 5}
    """
    return set([n for n in bits if isinstance(n, int)])


def get_free_net(nets):
    """
    Given a set of used net indices, returns a new, free index.

    >>> get_free_net({0, 1, 2,    4, 5, 6})
    3
    >>> get_free_net({0, 1, 2, 3, 4, 5, 6})
    7
    """
    sorted_nets = sorted(list(nets))

    # Find a gap in the sequence
    for i in range(len(nets) - 1):
        n0 = sorted_nets[i]
        n1 = sorted_nets[i + 1]
        if n1 != (n0 + 1):
            return n0 + 1

    # No gap was found, return max + 1.
    return sorted_nets[-1] + 1


# =============================================================================


def find_and_split_inout_ports_and_nets(design, module_name):
    """
    This function takes a given module from the design and splits all of its
    inout ports into pairs of inputs and outputs. Newly created ports are
    given suffixed. For example an inout port named "A" is going to be replaced
    by a pair consisting of "A_$inp" and "A_$out" ports.

    The function also looks for "netnames" that correspond to inout ports being
    split. These ones are removed and replaced with new ones related to newly
    added input and output ports.

    If any other "netname" mentions a net index connected to a former inout
    port, then the index is removed from the "netname" (replaced by "x"). If
    there are only "x" left in the "netname", then it is removed.

    The function returns port name map and net index map. The port map is a
    list of pairs (input name, input/output name). There are two entries per
    an inout. The net map is a dict indexed by indices of nets associated with
    inout ports. Each item contains a dict like {"i": int, "o": int} with
    indices of the inout net split products.
    """

    # Get the module
    module = design["modules"][module_name]

    # Find indices of all used nets
    nets = set()
    for port in module["ports"].values():
        nets |= get_nets(port["bits"])

    for netname in module["netnames"].values():
        nets |= get_nets(netname["bits"])

    for cell in module["cells"].values():
        for connection in cell["connections"].values():
            nets |= get_nets(connection)

    # Get all inout ports
    inouts = {
        k: v
        for k, v in module["ports"].items()
        if v["direction"] == "inout"
    }

    # Split ports
    new_ports = {}
    net_map = {}
    port_map = []
    for name, port in inouts.items():

        # Remove the inout port from the module
        del module["ports"][name]
        nets -= get_nets(port["bits"])

        # Make an input and output port
        for dir in ["input", "output"]:
            new_name = name + "_$" + dir[:3]
            new_port = {"direction": dir, "bits": []}

            print("Mapping port '{}' to '{}'".format(name, new_name))

            for n in port["bits"]:
                if isinstance(n, int):
                    mapped_n = get_free_net(nets)
                    print("Mapping net {} to {} ({})".format(n, mapped_n, dir))

                    if n not in net_map:
                        net_map[n] = {}
                    net_map[n][dir[0]] = mapped_n
                    nets.add(mapped_n)

                    new_port["bits"].append(mapped_n)
                else:
                    new_port["bits"].append(n)

            port_map.append((
                name,
                new_name,
            ))
            new_ports[new_name] = new_port

    # Add inputs and outputs
    module["ports"].update(new_ports)

    # .....................................................
    netnames = module["netnames"]

    # Remove netnames related to inout ports
    for name, net in list(module["netnames"].items()):
        if name in inouts:
            print("Removing netname '{}'".format(name))
            del netnames[name]

    # Remove remapped nets
    for name, net in list(netnames.items()):

        # Remove "bits" used by the net that were re-mapped.
        if len(set(net["bits"]) & set(net_map.keys())):

            # Remove
            net["bits"] = ["x" if b in net_map else b for b in net["bits"]]

            # If there is nothing left, remove the whole net.
            if all([b == "x" for b in net["bits"]]):
                print("Removing netname '{}'".format(name))
                del netnames[name]

    # Add netnames related to new input and output ports
    for name, port in new_ports.items():
        netnames[name] = {
            "hide_name": 0,
            "bits": port["bits"],
            "attributes": {}
        }

    return port_map, net_map


def remap_connections(design, module_name, net_map):
    """
    This function remaps cell connections that mention inout ports being split.

    The function loops over all cells and their ports. If a port contains
    a connection to an inout net then the connection is remapped according to
    the given net_map. Only ports which names ends on "_$inp" and "_$out" are
    affected.
    """

    module = design["modules"][module_name]
    cells = module["cells"]

    # Process cells
    for name, cell in cells.items():
        if "port_directions" not in cell:
            continue

        port_directions = cell["port_directions"]
        connections = cell["connections"]

        # Process cell connections
        for port_name, port_nets in list(connections.items()):

            # Skip if no net of this connection were remapped
            if len(set(net_map.keys()) & set(port_nets)) == 0:
                continue

            # Remove connections to the output net from input port and vice
            # versa.
            for dir in ["input", "output"]:
                if port_directions[port_name] == dir and \
                   port_name.endswith("$" + dir[:3]):

                    for i, n in enumerate(port_nets):
                        if n in net_map:
                            mapped_n = net_map[n][dir[0]]
                            port_nets[i] = mapped_n
                            print(
                                "Mapping connection {}.{}[{}] from {} to {}".
                                format(name, port_name, i, n, mapped_n)
                            )


# =============================================================================


def main():
    """
    The main.
    """

    # Parse args
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("-i", required=True, type=str, help="Input JSON")
    parser.add_argument("-o", default=None, type=str, help="Output JSON")

    args = parser.parse_args()

    # Output name
    if args.o is None:
        args.o = os.path.splitext(args.i)[0] + "_out.json"

    # Read the design
    with open(args.i, "r") as fp:
        design = json.load(fp)

    # Find the top module
    top_name = find_top_module(design)

    # Find and split inouot ports
    port_map, net_map = find_and_split_inout_ports_and_nets(design, top_name)
    # Remap cell connections
    remap_connections(design, top_name, net_map)

    # Write the design
    with open(args.o, "w") as fp:
        json.dump(design, fp, sort_keys=True, indent=2)


if __name__ == "__main__":
    main()