"""
Helper functions for reading and writing hardware files.

Each of the functions in inputoutput take a block and a file descriptor.
The functions provided either read the file and update the Block
accordingly, or write information from the Block out to the file.
"""

from __future__ import print_function, unicode_literals
import re
import collections

from .pyrtlexceptions import PyrtlError, PyrtlInternalError
from .core import working_block, _NameSanitizer
from .wire import WireVector, Input, Output, Const, Register, next_tempvar_name
from .corecircuits import concat_list
from .memory import RomBlock
from .passes import two_way_concat, one_bit_selects


def _natural_sort_key(key):
    """ Convert the key into a form such that it will be sorted naturally,
        e.g. such that "tmp4" appears before "tmp18".
    """
    def convert(text):
        return int(text) if text.isdigit() else text
    return [convert(c) for c in re.split(r'(\d+)', key)]


def _net_sorted(logic, name_mapper=lambda w: w.name):
    # Sort nets based on the name of the destination
    # wire, unless it's a memory write net.
    def natural_keys(n):
        if n.op == '@':
            # Sort based on the name of the wr_en wire, since
            # this particular net is used within 'always begin ... end'
            # blocks for memory update logic.
            key = str(n.args[2])
        else:
            key = name_mapper(n.dests[0])
        return _natural_sort_key(key)
    return sorted(logic, key=natural_keys)


def _name_sorted(wires, name_mapper=lambda w: w.name):
    return sorted(wires, key=lambda w: _natural_sort_key(name_mapper(w)))

# -----------------------------------------------------------------
#     __          ___
#    |__) |    | |___
#    |__) |___ | |


class Subcircuit:
    """
    This is a way to create and track per-module-instance wire names, so there
    are not name clashes when we instantiate a module more than once.
    """

    def __init__(self, model, is_top=False, clk_set={'clk'}, block=None):
        self.model = model
        self.is_top = is_top
        self.clk_set = clk_set
        self.block = working_block(block)
        self.inputs = {}
        self.outputs = {}
        self.wirevector_by_name = {}

    def add_input(self, original_name, wire):
        self.inputs[original_name] = wire
        self.wirevector_by_name[original_name] = wire

    def add_output(self, original_name, wire):
        self.outputs[original_name] = wire
        self.wirevector_by_name[original_name] = wire

    def add_reg(self, original_name, wire):
        self.wirevector_by_name[original_name] = wire

    def add_clock(self, clock_name):
        self.clk_set.add(clock_name)

    def twire(self, x):
        """ Find or make wire named x and return it. """
        s = self.wirevector_by_name.get(x)
        if s is None:
            # Purposefully *not* setting its name to 'x', so we don't have name clashes
            s = WireVector(bitwidth=1)
            self.wirevector_by_name[x] = s
        return s


def input_from_blif(blif, block=None, merge_io_vectors=True, clock_name='clk', top_model=None):
    """ Read an open BLIF file or string as input, updating the block appropriately.

    :param blif: An open BLIF file to read
    :param block: The block where the logic will be added
    :param merge_io_vectors: If True, Input/Output wires whose names differ only
        by a indexing subscript (e.g. 1-bit wires 'a[0]' and 'a[1]') will be combined
        into a single Input/Output (e.g. a 2-bit wire 'a').
    :param clock_name: The name of the clock (defaults to 'clk')
    :param top_model: name of top-level model to instantiate; if None, defaults to first model
        listed in the BLIF

    If merge_io_vectors is True, then given 1-bit Input wires 'a[0]' and 'a[1]', these
    wires will be combined into a single 2-bit Input wire 'a' that can be accessed
    by name 'a' in the block. Otherwise if merge_io_vectors is False, the original 1-bit
    wires will be Input wires of the block. This holds similarly for Outputs.

    This assumes the following:
      * The BLIF has been flattened and there is only a single module
      * There is only one single shared clock and reset
      * Output is generated by Yosys with formals in a particular order

    It currently ignores the reset signal (which it assumes is input only to the flip flops).
    """
    import pyparsing
    import six
    from pyparsing import (Word, Literal, OneOrMore, ZeroOrMore,
                           Suppress, Group, Keyword, Optional, oneOf)

    block = working_block(block)

    try:
        blif_string = blif.read()
    except AttributeError:
        if isinstance(blif, six.string_types):
            blif_string = blif
        else:
            raise PyrtlError('input_from_blif expecting either open file or string')

    def SKeyword(x):
        return Suppress(Keyword(x))

    def SLiteral(x):
        return Suppress(Literal(x))

    # Begin BLIF language definition
    signal_start = pyparsing.alphas + '$:[]_<>\\/?'
    signal_middle = pyparsing.alphas + pyparsing.nums + '$:[]_<>\\/.?-'
    signal_id = Word(signal_start, signal_middle)
    header = SKeyword('.model') + signal_id('model_name')
    input_list = Group(SKeyword('.inputs') + ZeroOrMore(signal_id))('input_list')
    output_list = Group(SKeyword('.outputs') + ZeroOrMore(signal_id))('output_list')

    cover_atom = Word('01-')
    cover_list = Group(ZeroOrMore(cover_atom))('cover_list')
    namesignal_list = Group(OneOrMore(signal_id))('namesignal_list')
    name_def = Group(SKeyword('.names') + namesignal_list + cover_list)('name_def')

    # asynchronous Flip-flop
    dffas_formal = (SLiteral('C=') + signal_id('C')
                    + SLiteral('R=') + signal_id('R')
                    + SLiteral('D=') + signal_id('D')
                    + SLiteral('Q=') + signal_id('Q'))
    dffas_keyword = SKeyword('$_DFF_PN0_') | SKeyword('$_DFF_PP0_')
    dffas_def = Group(SKeyword('.subckt') + dffas_keyword + dffas_formal)('dffas_def')

    # synchronous Flip-flop
    dffs_init_val = Optional(oneOf("0 1 2 3"), default=Literal("0"))
    # TODO I think <type> and <control> ('re' and 'C') below are technically optional too
    dffs_def = Group(SKeyword('.latch')
                     + signal_id('D')
                     + signal_id('Q')
                     + SLiteral('re')
                     + signal_id('C')
                     + dffs_init_val('I'))('dffs_def')

    # model reference
    formal_actual = Group(signal_id('formal') + SLiteral('=')
                          + signal_id('actual'))('formal_actual')
    formal_actual_list = Group(OneOrMore(formal_actual))('formal_actual_list')
    model_name = signal_id('model_name')
    model_ref = Group(SKeyword('.subckt') + model_name + formal_actual_list)('model_ref')

    command_def = name_def | dffas_def | dffs_def | model_ref
    command_list = Group(OneOrMore(command_def))('command_list')

    footer = SKeyword('.end')
    model_def = Group(header + input_list + output_list + command_list + footer)
    model_list = OneOrMore(model_def)
    parser = model_list.ignore(pyparsing.pythonStyleComment)

    # Begin actually reading and parsing the BLIF file
    result = parser.parseString(blif_string, parseAll=True)
    ff_clk_set = set([])
    models = {}  # model name -> model, for subckt instantiation

    def extract_inputs(subckt):
        if subckt.is_top:
            # NOTE: Assumes that:
            # - Top-level inputs starting with the same prefix are part of the same wire
            # - Indices start at 0
            start_names = [re.sub(r'\[([0-9]+)\]$', '', x) for x in subckt.model['input_list']]
            name_counts = collections.Counter(start_names)
            for input_name in name_counts:
                bitwidth = name_counts[input_name]
                if input_name in subckt.clk_set:
                    continue
                elif bitwidth == 1:
                    wire_in = Input(bitwidth=1, name=input_name, block=block)
                    subckt.add_input(input_name, wire_in)
                    block.add_wirevector(wire_in)
                elif merge_io_vectors:
                    wire_in = Input(bitwidth=bitwidth, name=input_name, block=block)
                    for i in range(bitwidth):
                        bit_name = input_name + '[' + str(i) + ']'
                        bit_wire = WireVector(bitwidth=1, block=block)
                        bit_wire <<= wire_in[i]
                        subckt.add_input(bit_name, bit_wire)
                else:
                    for i in range(bitwidth):
                        bit_name = input_name + '[' + str(i) + ']'
                        wire_in = Input(bitwidth=1, name=bit_name, block=block)
                        subckt.add_input(bit_name, wire_in)
                        block.add_wirevector(wire_in)
        else:
            # For subckts:
            # - Never merge input vectors
            # - All inputs are 1-bit
            for input_name in subckt.model['input_list']:
                if input_name in subckt.clk_set:
                    continue
                wire_in = WireVector(bitwidth=1, block=block)  # Internal name prevents name clash
                subckt.add_input(input_name, wire_in)
                block.add_wirevector(wire_in)

    def extract_outputs(subckt):
        if subckt.is_top:
            # NOTE: Assumes that:
            # - Top-level outputs starting with the same prefix are part of the same wire
            # - Indices start at 0
            start_names = [re.sub(r'\[([0-9]+)\]$', '', x) for x in subckt.model['output_list']]
            name_counts = collections.Counter(start_names)
            for output_name in name_counts:
                bitwidth = name_counts[output_name]
                # To allow an output wire to be used as an argument (legal in BLIF),
                # we need to create an intermediate wire, which will be used in twire()
                # whenever the original wire is referenced. For example, given 2-bit Output 'a',
                # every access to 'a[1]' will really be a reference to 'a[1]_i', a normal
                # WireVector connected to 'a[1]'. A key property is that the name by
                # which all other parts of the code refer to this wire doesn't change;
                # the only thing that changes is what underlying wire is used.
                if bitwidth == 1:
                    bit_internal = WireVector(bitwidth=1, block=block)
                    bit_out = Output(bitwidth=1, name=output_name, block=block)
                    bit_out <<= bit_internal
                    # NOTE this is important: redirecting user-visible name to internal wire
                    subckt.add_output(output_name, bit_internal)
                elif merge_io_vectors:
                    wire_out = Output(bitwidth=bitwidth, name=output_name, block=block)
                    bit_list = []
                    for i in range(bitwidth):
                        bit_name = output_name + '[' + str(i) + ']'
                        bit_wire = WireVector(bitwidth=1, block=block)
                        bit_list.append(bit_wire)
                        subckt.add_output(bit_name, bit_wire)
                    wire_out <<= concat_list(bit_list)
                else:
                    for i in range(bitwidth):
                        bit_name = output_name + '[' + str(i) + ']'
                        bit_internal = WireVector(bitwidth=1, block=block)
                        bit_out = Output(bitwidth=1, name=bit_name, block=block)
                        bit_out <<= bit_internal
                        # NOTE this is important: redirecting user-visible name to internal wire
                        subckt.add_output(bit_name, bit_internal)
        else:
            # For subckts:
            # - Never merge outputs vectors
            # - All outputs are 1-bit
            for output_name in subckt.model['output_list']:
                bit_out = WireVector(bitwidth=1, block=block)
                block.add_wirevector(bit_out)
                subckt.add_output(output_name, bit_out)

    def extract_commands(subckt):
        # for each "command" (dff or net) in the model
        for command in subckt.model['command_list']:
            # if it is a net (specified as a cover)
            if command.getName() == 'name_def':
                extract_cover(subckt, command)
            # else if the command is a d flop flop
            elif command.getName() == 'dffas_def' or command.getName() == 'dffs_def':
                extract_flop(subckt, command)
            elif command.getName() == 'model_ref':
                extract_model_reference(subckt, command)
            else:
                raise PyrtlError('unknown command type')

    def extract_cover(subckt, command):

        def twire(w):
            return subckt.twire(w)

        # pylint: disable=invalid-unary-operand-type
        netio = command['namesignal_list']
        if len(command['cover_list']) == 0:
            output_wire = twire(netio[0])
            output_wire <<= Const(0, bitwidth=1, block=block)  # const "FALSE"
        elif command['cover_list'].asList() == ['1']:
            output_wire = twire(netio[0])
            output_wire <<= Const(1, bitwidth=1, block=block)  # const "TRUE"
        elif command['cover_list'].asList() == ['1', '1']:
            # Populate clock list if one input is already a clock
            if (netio[1] in subckt.clk_set):
                subckt.add_clock(netio[0])
            elif (netio[0] in subckt.clk_set):
                subckt.add_clock(netio[1])
            else:
                output_wire = twire(netio[1])
                output_wire <<= twire(netio[0])  # simple wire
        elif command['cover_list'].asList() == ['0', '1']:
            output_wire = twire(netio[1])
            output_wire <<= ~ twire(netio[0])  # not gate
        elif command['cover_list'].asList() == ['11', '1']:
            output_wire = twire(netio[2])
            output_wire <<= twire(netio[0]) & twire(netio[1])  # and gate
        elif command['cover_list'].asList() == ['00', '1']:
            output_wire = twire(netio[2])
            output_wire <<= ~ (twire(netio[0]) | twire(netio[1]))  # nor gate
        elif command['cover_list'].asList() == ['1-', '1', '-1', '1']:
            output_wire = twire(netio[2])
            output_wire <<= twire(netio[0]) | twire(netio[1])  # or gate
        elif command['cover_list'].asList() == ['10', '1', '01', '1']:
            output_wire = twire(netio[2])
            output_wire <<= twire(netio[0]) ^ twire(netio[1])  # xor gate
        elif command['cover_list'].asList() == ['1-0', '1', '-11', '1']:
            output_wire = twire(netio[3])
            output_wire <<= (twire(netio[0]) & ~ twire(netio[2])) \
                | (twire(netio[1]) & twire(netio[2]))   # mux
        elif command['cover_list'].asList() == ['-00', '1', '0-0', '1']:
            output_wire = twire(netio[3])
            output_wire <<= (~twire(netio[1]) & ~twire(netio[2])) \
                | (~twire(netio[0]) & ~twire(netio[2]))
        else:
            raise PyrtlError('Blif file with unknown logic cover set "%s"'
                             '(currently gates are hard coded)' % command['cover_list'])

    def extract_flop(subckt, command):

        def twire(w):
            return subckt.twire(w)

        if(command['C'] not in ff_clk_set):
            ff_clk_set.add(command['C'])

        # Create register and assign next state to D and output to Q
        regname = command['Q'] + '_reg'
        flop = Register(bitwidth=1)
        subckt.add_reg(regname, flop)
        flop.next <<= twire(command['D'])
        flop_output = twire(command['Q'])
        init_val = command['I']
        if init_val == "1":
            # e.g. in Verilog: `initial reg <= 1;`
            raise PyrtlError("Initializing latches to 1 is not supported. "
                             "Acceptable values are: 0, 2 (don't care), and 3 (unknown); "
                             "in any case, PyRTL will ensure all stateful elements come up 0. "
                             "For finer control over the initial value, use specialized reset "
                             "logic.")
        flop_output <<= flop

    def extract_model_reference(parent, command):

        def twire(w):
            return parent.twire(w)

        def get_formal_connected_to_parent_clocks():
            clks = set()
            for fa in command['formal_actual_list']:
                if fa['actual'] in parent.clk_set:
                    clks.add(fa['formal'])
            return clks
        formal_clks = get_formal_connected_to_parent_clocks()

        subckt = Subcircuit(models[command['model_name']], clk_set=formal_clks, block=block)
        instantiate(subckt)
        for fa in command['formal_actual_list']:
            formal = fa['formal']
            actual = fa['actual']
            if actual in parent.clk_set:
                assert(formal in subckt.clk_set)
                # We didn't create an input wire corresponding to this.
                continue
            elif formal in subckt.inputs:
                wf = subckt.inputs[formal]
                wa = twire(actual)
                wf <<= wa
            elif formal in subckt.outputs:
                wf = subckt.outputs[formal]
                wa = twire(actual)
                wa <<= wf
            else:
                raise PyrtlError("%s formal parameter is neither an input nor output of subckt %s"
                                 % (formal, command['model_name']))

    def instantiate(subckt):
        extract_inputs(subckt)
        extract_outputs(subckt)
        extract_commands(subckt)

    # Get all model definitions
    for model in result:
        if not top_model:
            top_model = model['model_name']
        models[model['model_name']] = model

    top = Subcircuit(models[top_model], is_top=True, clk_set={clock_name}, block=block)
    instantiate(top)


# ----------------------------------------------------------------
#         ___  __          __   __
#   \  / |__  |__) | |    /  \ / _`
#    \/  |___ |  \ | |___ \__/ \__>
#


def output_to_verilog(dest_file, block=None):
    """ A function to walk the block and output it in verilog format to the open file. """

    block = working_block(block)
    file = dest_file
    internal_names = _VerilogSanitizer('_ver_out_tmp_')

    for wire in block.wirevector_set:
        internal_names.make_valid_string(wire.name)

    def varname(wire):
        return internal_names[wire.name]

    _to_verilog_header(file, block, varname)
    _to_verilog_combinational(file, block, varname)
    _to_verilog_sequential(file, block, varname)
    _to_verilog_memories(file, block, varname)
    _to_verilog_footer(file)


def OutputToVerilog(dest_file, block=None):
    """ A deprecated function to output verilog, use "output_to_verilog" instead. """
    return output_to_verilog(dest_file, block)


class _VerilogSanitizer(_NameSanitizer):
    _ver_regex = r'[_A-Za-z][_a-zA-Z0-9\$]*$'

    _verilog_reserved = \
        """always and assign automatic begin buf bufif0 bufif1 case casex casez cell cmos
        config deassign default defparam design disable edge else end endcase endconfig
        endfunction endgenerate endmodule endprimitive endspecify endtable endtask
        event for force forever fork function generate genvar highz0 highz1 if ifnone
        incdir include initial inout input instance integer join large liblist library
        localparam macromodule medium module nand negedge nmos nor noshowcancelledno
        not notif0 notif1 or output parameter pmos posedge primitive pull0 pull1
        pulldown pullup pulsestyle_oneventglitch pulsestyle_ondetectglitch remos real
        realtime reg release repeat rnmos rpmos rtran rtranif0 rtranif1 scalared
        showcancelled signed small specify specparam strong0 strong1 supply0 supply1
        table task time tran tranif0 tranif1 tri tri0 tri1 triand trior trireg unsigned
        use vectored wait wand weak0 weak1 while wire wor xnor xor
        """

    def __init__(self, internal_prefix='_sani_temp', map_valid_vals=True):
        self._verilog_reserved_set = frozenset(self._verilog_reserved.split())
        super(_VerilogSanitizer, self).__init__(self._ver_regex, internal_prefix,
                                                map_valid_vals, self._extra_checks)

    def _extra_checks(self, str):
        return(str not in self._verilog_reserved_set  # is not a Verilog reserved keyword
               and str != 'clk'                       # not the clock signal
               and len(str) <= 1024)                  # not too long to be a Verilog id


def _verilog_vector_size_decl(n):
    return '' if n == 1 else '[{:d}:0]'.format(n - 1)


def _verilog_vector_decl(w):
    return _verilog_vector_size_decl(len(w))


def _verilog_block_parts(block):
    inputs = block.wirevector_subset(Input)
    outputs = block.wirevector_subset(Output)
    registers = block.wirevector_subset(Register)
    wires = block.wirevector_subset() - (inputs | outputs | registers)
    memories = {n.op_param[1] for n in block.logic_subset('m@')}
    return inputs, outputs, registers, wires, memories


def _to_verilog_header(file, block, varname):
    """ Print the header of the verilog implementation. """

    def name_sorted(wires):
        return _name_sorted(wires, name_mapper=varname)

    def name_list(wires):
        return [varname(w) for w in wires]

    print('// Generated automatically via PyRTL', file=file)
    print('// As one initial test of synthesis, map to FPGA with:', file=file)
    print('//   yosys -p "synth_xilinx -top toplevel" thisfile.v\n', file=file)

    inputs, outputs, registers, wires, memories = _verilog_block_parts(block)

    # module name
    io_list = ['clk'] + name_list(name_sorted(inputs)) + name_list(name_sorted(outputs))
    if any(w.startswith('tmp') for w in io_list):
        raise PyrtlError('input or output with name starting with "tmp" indicates unnamed IO')
    io_list_str = ', '.join(io_list)
    print('module toplevel({:s});'.format(io_list_str), file=file)

    # inputs and outputs
    print('    input clk;', file=file)
    for w in name_sorted(inputs):
        print('    input{:s} {:s};'.format(_verilog_vector_decl(w), varname(w)), file=file)
    for w in name_sorted(outputs):
        print('    output{:s} {:s};'.format(_verilog_vector_decl(w), varname(w)), file=file)
    print('', file=file)

    # memories and registers
    for m in sorted(memories, key=lambda m: m.id):
        memwidth_str = _verilog_vector_size_decl(m.bitwidth)
        memsize_str = _verilog_vector_size_decl(1 << m.addrwidth)
        print('    reg{:s} mem_{}{:s}; //{}'.format(memwidth_str, m.id,
                                                    memsize_str, m.name), file=file)
    for w in name_sorted(registers):
        print('    reg{:s} {:s};'.format(_verilog_vector_decl(w), varname(w)), file=file)
    if (memories or registers):
        print('', file=file)

    # wires
    for w in name_sorted(wires):
        print('    wire{:s} {:s};'.format(_verilog_vector_decl(w), varname(w)), file=file)
    print('', file=file)

    # Write the initial values for read-only memories.
    # If we ever add support outside of simulation for initial values
    #  for MemBlocks, that would also go here.
    roms = {m for m in memories if isinstance(m, RomBlock)}
    for m in sorted(roms, key=lambda m: m.id):
        print('    initial begin', file=file)
        for i in range(1 << m.addrwidth):
            mem_elem_str = 'mem_{}[{:d}]'.format(m.id, i)
            mem_data_str = "{:d}'h{:x}".format(m.bitwidth, m._get_read_data(i))
            print('        {:s}={:s};'.format(mem_elem_str, mem_data_str), file=file)
        print('    end', file=file)
        print('', file=file)


def _to_verilog_combinational(file, block, varname):
    """ Print the combinational logic of the verilog implementation. """

    def name_sorted(wires):
        return _name_sorted(wires, name_mapper=varname)

    print('    // Combinational', file=file)

    # assign constants (these could be folded for readability later)
    for const in name_sorted(block.wirevector_subset(Const)):
        print('    assign {:s} = {:d};'.format(varname(const), const.val), file=file)

    # walk the block and output combination logic
    for net in _net_sorted(block.logic, varname):
        if net.op in 'w~':  # unary ops
            opstr = '' if net.op == 'w' else net.op
            t = (varname(net.dests[0]), opstr, varname(net.args[0]))
            print('    assign %s = %s%s;' % t, file=file)
        elif net.op in '&|^+-*<>':  # binary ops
            t = (varname(net.dests[0]), varname(net.args[0]),
                 net.op, varname(net.args[1]))
            print('    assign %s = %s %s %s;' % t, file=file)
        elif net.op == '=':
            t = (varname(net.dests[0]), varname(net.args[0]),
                 varname(net.args[1]))
            print('    assign %s = %s == %s;' % t, file=file)
        elif net.op == 'x':
            # note that the argument order for 'x' is backwards from the ternary operator
            t = (varname(net.dests[0]), varname(net.args[0]),
                 varname(net.args[2]), varname(net.args[1]))
            print('    assign %s = %s ? %s : %s;' % t, file=file)
        elif net.op == 'c':
            catlist = ', '.join([varname(w) for w in net.args])
            t = (varname(net.dests[0]), catlist)
            print('    assign %s = {%s};' % t, file=file)
        elif net.op == 's':
            # someone please check if we need this special handling for scalars
            catlist = ', '.join([varname(net.args[0]) + '[%s]' % str(i)
                                if len(net.args[0]) > 1 else varname(net.args[0])
                                for i in reversed(net.op_param)])
            t = (varname(net.dests[0]), catlist)
            print('    assign %s = {%s};' % t, file=file)
        elif net.op in 'rm@':
            pass  # do nothing for registers and memories
        else:
            raise PyrtlInternalError("nets with op '{}' not supported".format(net.op))
    print('', file=file)


def _to_verilog_sequential(file, block, varname):
    """ Print the sequential logic of the verilog implementation. """
    if not block.logic_subset(op='r'):
        return

    print('    // Registers', file=file)
    print('    always @( posedge clk )', file=file)
    print('    begin', file=file)
    for net in _net_sorted(block.logic, varname):
        if net.op == 'r':
            dest, src = (varname(net.dests[0]), varname(net.args[0]))
            print('        {:s} <= {:s};'.format(dest, src), file=file)
    print('    end', file=file)
    print('', file=file)


def _to_verilog_memories(file, block, varname):
    """ Print the memories of the verilog implementation. """
    memories = {n.op_param[1] for n in block.logic_subset('m@')}
    for m in sorted(memories, key=lambda m: m.id):
        print('    // Memory mem_{}: {}'.format(m.id, m.name), file=file)
        print('    always @( posedge clk )', file=file)
        print('    begin', file=file)
        for net in _net_sorted(block.logic_subset('@'), varname):
            if net.op_param[1] == m:
                t = (varname(net.args[2]), net.op_param[0],
                     varname(net.args[0]), varname(net.args[1]))
                print(('        if (%s) begin\n'
                       '            mem_%s[%s] <= %s;\n'
                       '        end') % t, file=file)
        print('    end', file=file)
        for net in _net_sorted(block.logic_subset('m'), varname):
            if net.op_param[1] == m:
                dest = varname(net.dests[0])
                m_id = net.op_param[0]
                index = varname(net.args[0])
                print('    assign {:s} = mem_{}[{:s}];'.format(dest, m_id, index), file=file)
        print('', file=file)


def _to_verilog_footer(file):
    print('endmodule\n', file=file)


def output_verilog_testbench(dest_file, simulation_trace=None, toplevel_include=None,
                             vcd="waveform.vcd", cmd=None, block=None):
    """ Output a Verilog testbench for the block/inputs used in the simulation trace.

    :param dest_file: an open file to which the test bench will be printed.
    :param simulation_trace: a simulation trace from which the inputs will be extracted
        for inclusion in the test bench. The test bench generated will just replay the
        inputs played to the simulation cycle by cycle. The default values for all
        registers and memories will be based on the trace, otherwise they will be initialized
        to 0.
    :param toplevel_include: name of the file containing the toplevel module this testbench
        is testing. If not None, an `include` directive will be added to the top.
    :param vcd: By default the testbench generator will include a command in the testbench
        to write the output of the testbench execution to a .vcd file (via $dumpfile), and
        this parameter is the string of the name of the file to use. If None is specified
        instead, then no dumpfile will be used.
    :param cmd: The string passed as cmd will be copied verbatim into the testbench
        just before the end of each cycle. This is useful for doing things like printing
        specific values out during testbench evaluation (e.g. cmd='$display("%d", out);'
        will instruct the testbench to print the value of 'out' every cycle which can then
        be compared easy with a reference).

    The test bench does not return any values.

    Example 1 (writing testbench to a string)::

        with io.StringIO() as tbfile:
            pyrtl.output_verilog_testbench(dest_file=tbfile, simulation_trace=sim_trace)

    Example 2 (testbench in same file as verilog)::

        with open('hardware.v', 'w') as fp:
            output_to_verilog(fp)
            output_verilog_testbench(fp, sim.tracer, vcd=None, cmd='$display("%d", out);')

    """
    block = working_block(block)

    inputs, outputs, registers, wires, memories = _verilog_block_parts(block)

    ver_name = _VerilogSanitizer('_ver_out_tmp_')
    for wire in block.wirevector_set:
        ver_name.make_valid_string(wire.name)

    def name_sorted(wires):
        return _name_sorted(wires, name_mapper=lambda w: ver_name[w.name])

    def name_list(wires):
        return [ver_name[w.name] for w in wires]

    def init_regvalue(r):
        if simulation_trace:
            return simulation_trace.init_regvalue.get(r, simulation_trace.default_value)
        else:
            return 0

    def init_memvalue(m, ix):
        # Return None if not present, or if already equal to default value, so we know not to
        # emit any additional Verilog initing this mem address.
        if simulation_trace:
            if m not in simulation_trace.init_memvalue:
                return None
            v = simulation_trace.init_memvalue[m].get(ix, simulation_trace.default_value)
            return None if v == simulation_trace.default_value else v
        else:
            return None

    def default_value():
        return simulation_trace.default_value if simulation_trace else 0

    # Output an include, if given
    if toplevel_include:
        print('`include "{:s}"'.format(toplevel_include), file=dest_file)
        print('', file=dest_file)

    # Output header
    print('module tb();', file=dest_file)

    # Declare all block inputs as reg
    print('    reg clk;', file=dest_file)
    for w in name_sorted(inputs):
        print('    reg{:s} {:s};'.format(_verilog_vector_decl(w), ver_name[w.name]),
              file=dest_file)

    # Declare all block outputs as wires
    for w in name_sorted(outputs):
        print('    wire{:s} {:s};'.format(_verilog_vector_decl(w), ver_name[w.name]),
              file=dest_file)
    print('', file=dest_file)

    # Declare an integer used for init of memories
    print('    integer tb_iter;', file=dest_file)

    # Instantiate logic block
    io_list = ['clk'] + name_list(name_sorted(inputs)) + name_list(name_sorted(outputs))
    io_list_str = ['.{0:s}({0:s})'.format(w) for w in io_list]
    print('    toplevel block({:s});\n'.format(', '.join(io_list_str)), file=dest_file)

    # Generate clock signal
    print('    always', file=dest_file)
    print('        #5 clk = ~clk;\n', file=dest_file)

    # Move through all steps of trace, writing out input assignments per cycle
    print('    initial begin', file=dest_file)

    # If a VCD output is requested, set that up
    if vcd:
        print('        $dumpfile ("%s");' % vcd, file=dest_file)
        print('        $dumpvars;\n', file=dest_file)

    # Initialize clk, and all the registers and memories
    print('        clk = 0;', file=dest_file)
    for r in name_sorted(registers):
        print('        block.%s = %d;' % (ver_name[r.name], init_regvalue(r)), file=dest_file)
    for m in sorted(memories, key=lambda m: m.id):
        max_iter = 1 << m.addrwidth
        print('        for (tb_iter = 0; tb_iter < %d; tb_iter++) '
              'begin block.mem_%s[tb_iter] = %d; end' % (max_iter, m.id, default_value()),
              file=dest_file)
        for ix in range(max_iter):
            # Now just individually update the memory values that aren't the default
            val = init_memvalue(m.id, ix)
            if val is not None:
                print('        block.mem_%s[%d] = %d;' % (m.id, ix, val), file=dest_file)

    if simulation_trace:
        tracelen = max(len(t) for t in simulation_trace.trace.values())
        for i in range(tracelen):
            for w in name_sorted(inputs):
                print('        {:s} = {:s}{:d};'.format(
                    ver_name[w.name],
                    "{:d}'d".format(len(w)),
                    simulation_trace.trace[w][i]), file=dest_file)
            if cmd:
                print('        %s' % cmd, file=dest_file)
            print('\n        #10', file=dest_file)

    # Footer
    print('        $finish;', file=dest_file)
    print('    end', file=dest_file)
    print('endmodule', file=dest_file)


# ----------------------------------------------------------------
#    ___    __   __  ___
#   |___ | |__) |__)  |  |
#   |    | |  \ |  \  |  |___
#

def output_to_firrtl(open_file, rom_blocks=None, block=None):
    """ Output the block as FIRRTL code to the output file.

    :param open_file: File to write to
    :param rom_blocks: List of ROM blocks to be initialized
    :param block: Block to use (defaults to working block)

    If ROM is initialized in PyRTL code, you can pass in the rom_blocks as a list [rom1, rom2, ...].
    """
    block = working_block(block)

    # FIRRTL only allows 'bits' operations to have two parameters: a high and low
    # index representing the inclusive bounds of a contiguous range. PyRTL uses
    # slice syntax, which aren't always contiguous, so we need to convert them.
    one_bit_selects(block=block)  # pylint: disable=no-value-for-parameter,unexpected-keyword-arg

    # FIRRTL only allows 'concatenate' operations to have two arguments,
    # but PyRTL's 'c' op allows an arbitrary number of wires. We need to convert
    # these n-way concats to series of two-way concats accordingly.
    two_way_concat(block=block)  # pylint: disable=no-value-for-parameter,unexpected-keyword-arg

    f = open_file
    # write out all the implicit stuff
    f.write("circuit Example :\n")
    f.write("  module Example :\n")
    f.write("    input clock : Clock\n    input reset : UInt<1>\n")
    # write out IO signals, wires and registers
    for wire in _name_sorted(block.wirevector_subset(Input)):
        f.write("    input %s : UInt<%d>\n" % (wire.name, wire.bitwidth))
    for wire in _name_sorted(block.wirevector_subset(Output)):
        f.write("    output %s : UInt<%d>\n" % (wire.name, wire.bitwidth))
    for wire in _name_sorted(block.wirevector_subset(exclude=(Input, Output, Register, Const))):
        f.write("    wire %s : UInt<%d>\n" % (wire.name, wire.bitwidth))
    for wire in _name_sorted(block.wirevector_subset(Register)):
        f.write("    reg %s : UInt<%d>, clock\n" % (wire.name, wire.bitwidth))
    for wire in _name_sorted(block.wirevector_subset(Const)):
        # some const is in the form like const_0_1'b1, is this legal operation?
        wire.name = wire.name.split("'").pop(0)
        f.write("    node %s = UInt<%d>(%d)\n" % (wire.name, wire.bitwidth, wire.val))
    f.write("\n")

    # write "Main"
    node_cntr = 0
    initializedMem = []
    for log_net in _net_sorted(block.logic_subset()):
        if log_net.op == '&':
            f.write("    %s <= and(%s, %s)\n" % (log_net.dests[0].name, log_net.args[0].name,
                                                 log_net.args[1].name))
        elif log_net.op == '|':
            f.write("    %s <= or(%s, %s)\n" % (log_net.dests[0].name, log_net.args[0].name,
                                                log_net.args[1].name))
        elif log_net.op == '^':
            f.write("    %s <= xor(%s, %s)\n" % (log_net.dests[0].name, log_net.args[0].name,
                                                 log_net.args[1].name))
        elif log_net.op == 'n':
            f.write("    node T_%d = and(%s, %s)\n" % (node_cntr, log_net.args[0].name,
                                                       log_net.args[1].name))
            f.write("    %s <= not(T_%d)\n" % (log_net.dests[0].name, node_cntr))
            node_cntr += 1
        elif log_net.op == '~':
            f.write("    %s <= not(%s)\n" % (log_net.dests[0].name, log_net.args[0].name))
        elif log_net.op == '+':
            f.write("    %s <= add(%s, %s)\n" % (log_net.dests[0].name, log_net.args[0].name,
                                                 log_net.args[1].name))
        elif log_net.op == '-':
            f.write("    %s <= sub(%s, %s)\n" % (log_net.dests[0].name, log_net.args[0].name,
                                                 log_net.args[1].name))
        elif log_net.op == '*':
            f.write("    %s <= mul(%s, %s)\n" % (log_net.dests[0].name, log_net.args[0].name,
                                                 log_net.args[1].name))
        elif log_net.op == '=':
            f.write("    %s <= eq(%s, %s)\n" % (log_net.dests[0].name, log_net.args[0].name,
                                                log_net.args[1].name))
        elif log_net.op == '<':
            f.write("    %s <= lt(%s, %s)\n" % (log_net.dests[0].name, log_net.args[0].name,
                                                log_net.args[1].name))
        elif log_net.op == '>':
            f.write("    %s <= gt(%s, %s)\n" % (log_net.dests[0].name, log_net.args[0].name,
                                                log_net.args[1].name))
        elif log_net.op == 'w':
            f.write("    %s <= %s\n" % (log_net.dests[0].name, log_net.args[0].name))
        elif log_net.op == 'x':
            f.write("    %s <= mux(%s, %s, %s)\n" % (log_net.dests[0].name, log_net.args[0].name,
                                                     log_net.args[2].name, log_net.args[1].name))
        elif log_net.op == 'c':
            if len(log_net.args) != 2:
                raise PyrtlInternalError(
                    "Expected concat net to have only two "
                    "argument wires; has %d" % len(log_net.args)
                )
            f.write("    %s <= cat(%s, %s)\n" % (log_net.dests[0].name, log_net.args[0].name,
                                                 log_net.args[1].name))
        elif log_net.op == 's':
            if len(log_net.op_param) != 1:
                raise PyrtlInternalError(
                    "Expected select net to have single "
                    "select bit; has %d" % len(log_net.op_param)
                )
            f.write("    %s <= bits(%s, %s, %s)\n" % (log_net.dests[0].name, log_net.args[0].name,
                                                      log_net.op_param[0], log_net.op_param[0]))
        elif log_net.op == 'r':
            f.write("    %s <= mux(reset, UInt<%s>(0), %s)\n" %
                    (log_net.dests[0].name, log_net.dests[0].bitwidth, log_net.args[0].name))
        elif log_net.op == 'm':
            # if there are rom blocks, need to be initialized
            if rom_blocks is not None:
                if not log_net.op_param[0] in initializedMem:
                    initializedMem.append(log_net.op_param[0])

                    # find corresponding rom block according to memid
                    curr_rom = next((x for x in rom_blocks if x.id == log_net.op_param[0]), None)
                    f.write("    wire %s : UInt<%s>[%s]\n" %
                            (log_net.op_param[1].name, log_net.op_param[1].bitwidth,
                             2**log_net.op_param[1].addrwidth))

                    # if rom data is a function, calculate the data first
                    if callable(curr_rom.data):
                        romdata = [curr_rom.data(i) for i in range(2**curr_rom.addrwidth)]
                        curr_rom.data = romdata

                    # write rom block initialization data
                    for i in range(len(curr_rom.data)):
                        f.write("    %s[%s] <= UInt<%s>(%s)\n" %
                                (log_net.op_param[1].name, i, log_net.op_param[1].bitwidth,
                                 curr_rom.data[i]))

                # write the connection
                f.write("    %s <= %s[%s]\n" % (log_net.dests[0].name, log_net.op_param[1].name,
                                                log_net.args[0].name))

            else:
                if not log_net.op_param[0] in initializedMem:
                    initializedMem.append(log_net.op_param[0])
                    f.write("    cmem %s_%s : UInt<%s>[%s]\n" %
                            (log_net.op_param[1].name, log_net.op_param[0],
                             log_net.op_param[1].bitwidth, 2**log_net.op_param[1].addrwidth))
                f.write("    infer mport T_%d  = %s_%s[%s], clock\n" %
                        (node_cntr, log_net.op_param[1].name, log_net.op_param[0],
                         log_net.args[0].name))
                f.write("    %s <= T_%d\n" % (log_net.dests[0].name, node_cntr))
                node_cntr += 1
        elif log_net.op == '@':
            if not log_net.op_param[0] in initializedMem:
                initializedMem.append(log_net.op_param[0])
                f.write("    cmem %s_%s : UInt<%s>[%s]\n" %
                        (log_net.op_param[1].name, log_net.op_param[0],
                         log_net.op_param[1].bitwidth, 2**log_net.op_param[1].addrwidth))
            f.write("    when %s :\n" % log_net.args[2].name)
            f.write("      infer mport T_%d  = %s_%s[%s], clock\n" %
                    (node_cntr, log_net.op_param[1].name, log_net.op_param[0],
                     log_net.args[0].name))
            f.write("      T_%d <= %s\n" % (node_cntr, log_net.args[1].name))
            f.write("      skip\n")
            node_cntr += 1
        else:
            pass

    return 0

# -----------------------------------------------------------------
#       __   __       __     __   ___       __
#    | /__` /    /\  /__`   |__) |__  |\ | /   |__|
#    | .__/ \__ /~~\ .__/   |__) |___ | \| \__ |  |


def input_from_iscas_bench(bench, block=None):
    ''' Import an ISCAS .bench file

    :param file bench: an open ISCAS .bench file to read
    :param block: block to add the imported logic (defaults to current working block)
    '''

    import pyparsing
    import six
    from pyparsing import (Word, Literal, OneOrMore, ZeroOrMore, Suppress, Group, Keyword, oneOf)

    block = working_block(block)

    try:
        bench_string = bench.read()
    except AttributeError:
        if isinstance(bench, six.string_types):
            bench_string = bench
        else:
            raise PyrtlError('input_from_bench expecting either open file or string')

    def SKeyword(x):
        return Suppress(Keyword(x))

    def SLiteral(x):
        return Suppress(Literal(x))

    # NOTE: The acceptable signal characters are based on viewing the I/O names
    # in the available ISCAS benchmark files, and may not be complete.
    signal_start = pyparsing.alphas + pyparsing.nums + '$:[]_<>\\/?'
    signal_middle = pyparsing.alphas + pyparsing.nums + '$:[]_<>\\/.?-'
    signal_id = Word(signal_start, signal_middle)

    gate_names = "AND OR NAND NOR XOR NOT BUFF DFF"

    src_list = Group(signal_id + ZeroOrMore(SLiteral(",") + signal_id))("src_list")
    net_def = Group(signal_id("dst") + SLiteral("=") + oneOf(gate_names)("gate")
                    + SLiteral("(") + src_list + SLiteral(")"))("net_def")

    input_def = Group(SKeyword("INPUT") + SLiteral("(")
                      + signal_id + SLiteral(")"))("input_def")
    output_def = Group(SKeyword("OUTPUT") + SLiteral("(")
                       + signal_id + SLiteral(")"))("output_def")
    command_def = input_def | output_def | net_def

    commands = OneOrMore(command_def)("command_list")
    parser = commands.ignore(pyparsing.pythonStyleComment)

    # Begin actually reading and parsing the BENCH file
    result = parser.parseString(bench_string, parseAll=True)

    output_to_internal = {}  # dict: name -> wire

    def twire(name):
        """ Find or make wire named 'name' and return it. """
        w = output_to_internal.get(name)
        if w is None:
            w = block.wirevector_by_name.get(name)
            if w is None:
                w = WireVector(bitwidth=1, name=name)
        return w

    for cmd in result["command_list"]:
        if cmd.getName() == "input_def":
            _wire_in = Input(bitwidth=1, name=str(cmd[0]), block=block)
        elif cmd.getName() == "output_def":
            # Create internal wire for indirection, since Outputs can't be inputs to nets in PyRTL
            wire_internal = WireVector(bitwidth=1, block=block)
            wire_out = Output(bitwidth=1, name=str(cmd[0]), block=block)
            wire_out <<= wire_internal
            output_to_internal[cmd[0]] = wire_internal
        elif cmd.getName() == "net_def":
            srcs = cmd["src_list"]
            if cmd["gate"] == "AND":
                dst_wire = twire(cmd["dst"])
                dst_wire <<= twire(srcs[0]) & twire(srcs[1])
            elif cmd["gate"] == "OR":
                dst_wire = twire(cmd["dst"])
                dst_wire <<= twire(srcs[0]) | twire(srcs[1])
            elif cmd["gate"] == "NAND":
                dst_wire = twire(cmd["dst"])
                dst_wire <<= twire(srcs[0]).nand(twire(srcs[1]))
            elif cmd["gate"] == "NOR":
                dst_wire = twire(cmd["dst"])
                dst_wire <<= ~(twire(srcs[0]) | twire(srcs[1]))
            elif cmd["gate"] == "XOR":
                dst_wire = twire(cmd["dst"])
                dst_wire <<= twire(srcs[0]) ^ twire(srcs[1])
            elif cmd["gate"] == "NOT":
                dst_wire = twire(cmd["dst"])
                dst_wire <<= ~twire(srcs[0])
            elif cmd["gate"] == "BUFF":
                dst_wire = twire(cmd["dst"])
                dst_wire <<= twire(srcs[0])
            elif cmd["gate"] == "DFF":
                dst_wire = twire(cmd["dst"])
                reg = Register(bitwidth=1)
                reg.next <<= twire(srcs[0])
                dst_wire <<= reg
            else:
                raise PyrtlError("Unexpected gate {%s}" % cmd["gate"])

    # Benchmarks like c1196, b18, etc. have inputs and outputs by the
    # same name, that are therefore directly connected. This pass will
    # rename the outputs so that this is still okay.
    for o in block.wirevector_subset(Output):
        inputs = [i for i in block.wirevector_subset(Input) if i.name == o.name]
        if inputs:
            if len(inputs) > 1:
                raise PyrtlError("More than one input found with the name %s" % inputs[0].name)
            i = inputs[0]
            o_internal = twire(o.name)
            o_internal <<= i
            o.name = next_tempvar_name()
            # Ensure the input is the one mapped by the original name
            block.wirevector_by_name[i.name] = i
            print("Found input and output wires with the same name. "
                  "Output '%s' has now been renamed to '%s'." % (i.name, o.name))
