from functools import reduce
import random
import visual


######## Width and height #######
N = 3
#M = 3
#################################



EMPTY = -1
WIRE = 0
SOURCE = 1
SINK = 2
types = ["Wire","Source","Sink"]

NORTH,EAST,SOUTH,WEST = range(4)
edges = ["North","East","South","West"]

directions = [(0,1),(1,0),(0,-1),(-1,0)]

def isWireConnected(puzzle,x,y):
    _,ds = puzzle[y][x]
    drs= [directions[d] for d in ds]
    for d in ds:
        dx,dy = directions[d]
        if 0 <= y+dy < len(puzzle) and 0 <= x+dx < len(puzzle[y+dy]):
            t2,ds2 = puzzle[y+dy][x+dx]
            if t2 == EMPTY:
                continue
            if (d+2)%4 not in ds2:
                return False
        else:
            return False
    return True

def gettile():
    while True:
        tile = (WIRE, [i for i in range(4) if random.random() < 0.4])
        if len(tile[1]) == 1:
            if random.random() < 0.1:
                return (random.choice([SINK,SOURCE]), tile[1])
        elif len(tile[1]) == 0:
            if random.random() < 0.2:
                return tile
        else:
            return tile
    

def consistent(puzzle,x,y):
    if not isWireConnected(puzzle,x,y):
        return False
    if y > 0:
        if not isWireConnected(puzzle,x,y-1):
            return False
    if x > 0:
        if not isWireConnected(puzzle,x-1,y):
            return False
    return True

def joined(puzzle,x1,y1,x2,y2):
    # is there a wire connecting the tiles at x1,y1 and x2,y2
    _,ds = puzzle[y1][x1]
    _,ds2 = puzzle[y2][x2]
    diff = (x2-x1,y2-y1)
    if diff not in [(0,1),(1,0),(0,-1),(-1,0)]:
        return False
    d = directions.index(diff)
    return (d+2)%4 in ds2 and d in ds

def neighbours(puzzle,x,y):
    _,ds = puzzle[y][x]
    for d in ds:
        dx,dy = directions[d]
        if 0 <= y+dy < len(puzzle) and 0 <= x+dx < len(puzzle[y+dy]):
            yield (x+dx,y+dy)

def reachable(puzzle,x,y):
    ttype,_ = puzzle[y][x]
    found = [(x,y)]
    ttype2 = SINK if ttype == SOURCE else SOURCE
    assert ttype in [SOURCE,SINK]
    
    flag = True
    while flag:
        flag = False
        for (x1,y1) in found.copy():
            for (x2,y2) in neighbours(puzzle,x1,y1):
                if (x2,y2) in found:
                    continue
                found.append((x2,y2))
                flag = True
                if puzzle[y2][x2][0] == ttype2:
                    return (x2,y2)
    return random.choice(found[1:])

seed1=0
def genpuzzle(w,h=None,seed=0):
    global seed1
    if h is None:
        h = w
    if seed == 0:
        seed = random.randint(0,9223372036854775807)
    random.seed(seed)
    seed1 = seed
    puzzle = [[(EMPTY,[]) for _ in range(w)] for _ in range(h)]
    for y in range(h):
        for x in range(w):
            puzzle[y][x] = gettile()
            while not consistent(puzzle,x,y):
                puzzle[y][x] = gettile()
    
    numsinks = (w+h)//5 + 1
    # add linked source to each sink and vice versa
    for y in range(h):
        for x in range(w):
            if puzzle[y][x][0] in [SOURCE,SINK]:
                x2,y2 = reachable(puzzle,x,y)
                puzzle[y2][x2] = (3-puzzle[y][x][0],puzzle[y2][x2][1])
    
    for _ in range(numsinks):
        x = random.randint(0,w-1)
        y = random.randint(0,h-1)
        if puzzle[y][x][1] != []:
            t1,t2 = random.choice([(SOURCE,SINK),(SINK,SOURCE)])
            puzzle[y][x] = (t1,puzzle[y][x][1])
            x2,y2 = reachable(puzzle,x,y)
            puzzle[y2][x2] = (t2,puzzle[y2][x2][1])
    
    return puzzle

def h(puzzle):
    output = "["
    for row in reversed(puzzle):
        ss = []
        for t in row:
            s = types[t[0]]# + " "
            s += "[" + ",".join(map(lambda x: edges[x], t[1])) + "]"
            ss.append(s)
        output += "[" + ",".join(ss) + "],"
    output = output[:-1] + "]"
    return output

def shuf(puzzle):
    def rot(rs,x):
        return (x[0], list(map(lambda y: (y+rs)%4, x[1])))
    return [[rot(random.randint(0,3),t) for t in ln] for ln in puzzle]


if __name__ == '__main__':
    try:
        p = genpuzzle(N,M)
    except NameError:
        p = genpuzzle(N)

    p_ = shuf(p)

    print(h(p))
    print('=======================================')
    print(h(p_))

    p2 = visual.join(p,p_)
    visual.show(p2)