from __future__ import print_function
from benchpress import util
import numpy as np
from bohrium.stdviews import D2P8, no_border

def wireworld_init(size):
    """TODO: Describe the data generated by this function."""

    data = np.zeros(
        (size*10+2, size*10+2),
        dtype   = np.uint8
    )
    data[1:-1,1:-1] = np.tile(np.array([
        [4,2,8,8,8,8,8,8,8,8],
        [8,0,0,0,8,0,0,0,0,0],
        [0,0,0,8,8,8,0,0,0,0],
        [8,0,0,0,8,0,0,0,0,0],
        [2,4,8,8,0,8,8,8,8,8],
        [4,2,8,8,8,8,8,8,8,8],
        [8,0,0,0,8,0,0,0,0,0],
        [0,0,0,8,8,8,0,0,0,0],
        [8,0,0,0,8,0,0,0,0,0],
        [2,4,8,8,0,8,8,8,8,8],
        ]),(size,size))

    return data

def wireworld(world, iterations):
    """TODO: Describe the benchmark."""

    sim     = no_border(world, 1)          # Active Machine
    stencil = D2P8(world)                  # Stencil for counting heads
    NC      = sum([v==2 for v in stencil]) # Count number of head neighbors
    for _ in range(iterations):
        # Mask conductor->head
        MASK = ((NC==1) & (sim==8)) | ((NC==2) & (sim==8))

        sim *= ~MASK                        # New head pos->0
        sim += np.array(MASK * 1, np.uint8) # New head pos->1
        MASK = (sim==8)                     # Mask non conductors
        sim *= ~MASK                        # conductors->0
        sim += np.array(MASK * 4, np.uint8) # conductors->4
        sim *= 2                            # Upgrade all to new state

        util.Benchmark().flush()

    return sim

def main():
    """
    Example parameter: --size=100*10.
    This will execute on a 1000x1000 dataset for 10 iterations.
    """
    B = util.Benchmark()
    (N, I) = B.size
    if B.inputfn:
        world = B.load_array()
    else:
        world = wireworld_init(N)

    if B.dumpinput:
        B.dump_arrays("wireworld", {"input": world})

    B.start()
    R = wireworld(world, I)
    B.stop()
    B.pprint()

    if B.outputfn:
        B.tofile(B.outputfn, {"res": R})

if __name__ == "__main__":
    main()
