from nmigen import *
from nmigen.asserts import *
from nmigen.test.utils import *
from ...sed.output_network import *

"""
Verification tasks for OutputNetwork
"""

class OutputNetworkSpec(Elaboratable):
    def elaborate(self, platform):
        m = Module()
        m.submodules.output_network = output_network = OutputNetwork(4, 2, [("data", 32), ("channel", 3)])

        # Model arbitrary inputs for network nodes
        for i in range(output_network.lane_count):
            m.d.comb += output_network.input[i].eq(AnySeq(1))
            m.d.comb += output_network.input[i].seqn.eq(AnySeq(output_network.seqn_width))
            m.d.comb += output_network.input[i].replace_occured.eq(0)
            m.d.comb += output_network.input[i].nondata_replace_occured.eq(0)
            for field, width in output_network.layout_payload:
                m.d.comb += getattr(output_network.input[i].payload, field).eq(AnySeq(width))
        # Indicator of when inputs from the first clock cycle make it
        # through the sorting network
        network_latency = latency(output_network.lane_count)
        counter = Signal(range(network_latency + 1))
        m.d.sync += counter.eq(counter + 1)
        with m.If(counter == network_latency):
            m.d.sync += counter.eq(counter)
        f_output_valid = Signal()
        m.d.comb += f_output_valid.eq(counter == network_latency)
        
        with m.If(f_output_valid):
            replacement_occurred = Signal()
            for node in output_network.output:
                with m.If(node.replace_occured):
                    m.d.comb += replacement_occurred.eq(1)
            channels_unique = Signal(reset=1)
            for node1 in range(len(output_network.input)):
                for node2 in range(node1):
                    k1 = Past(output_network.input[node1].payload.channel, clocks=network_latency)
                    k2 = Past(output_network.input[node2].payload.channel, clocks=network_latency)
                    with m.If(k1 == k2):
                        m.d.comb += channels_unique.eq(0)
            # If there are no replacements then:
            # - Input channel numbers are unique among (in)valid nodes
            # - All inputs make it through the sorting network unmodified
            # TODO
            # Otherwise, if there are replacements:
            # - There is a channel number collision among the valid or invalid
            #   nodes
            # - All channel numbers in valid inputs appear exactly once as a
            #   valid output
            # - All valid outputs correspond to a valid input modulo accounting
            #   information
            # - Channel numbers in invalid inputs not appearing in valid inputs
            #   as well never appear as a valid output
            # TODO

        return m

class OutputNetworkTestCase(FHDLTestCase):
    def verify(self):
        # Bounded model check
        self.assertFormal(
            OutputNetworkSpec(),
            mode="bmc", depth=40)
        # TODO: unbounded proof

OutputNetworkTestCase().verify()