import myhdl
from myhdl import *

from arith_utils import log2ceil
from arith_utils import SLOW, MEDIUM, FAST


def PrefixAnd(width, speed, PI, PO):

    n = width
    m = log2ceil(width)

    def fastPrefix():
        PT = Signal(intbv(0))

        @instance
        def comb():
            while 1:
                yield PI, PT
                PT.next[n:] = PI
                for l in range(1, m + 1):
                    for k in range(2 ** (m - l)):
                        for i in range(2 ** (l - 1)):
                            if (k * 2 ** l + i) < n:
                                PT.next[l * n + k * 2 ** l + i] = \
                                    PT[(l - 1) * n + k * 2 ** l + i]
                            if (k * 2 ** l + 2 ** (l - 1) + i) < n:
                                PT.next[l * n + k * 2 ** l + 2 ** (l - 1) + i] = \
                                    PT[(l - 1) * n + k * 2 ** l + 2 ** (l - 1) + i] & \
                                    PT[(l - 1) * n + k * 2 ** l + 2 ** (l - 1) - 1]
                PO.next = PT[(m + 1) * n:m * n]

        return comb

    def slowPrefix():
        PT = Signal(intbv(0))

        @instance
        def comb():
            while 1:
                yield PI, PT
                PT.next[0] = PI[0]
                for i in range(1, n):
                    PT.next[i] = PI[i] & PT[i - 1]
                PO.next = PT

        return comb

    if speed == SLOW:
        return slowPrefix()
    elif speed == FAST:
        return fastPrefix()
    else:
        raise NotImplementedError

