import numpy as np
from random import sample

def linear_to_pair(i, s):
    x = i // s
    y = i - x * s
    return x, y

def offset_inside(x, y, ox, oy, s):
    _x = (x + ox) % s
    _y = (y + oy) % s
    return (_x, _y)

def reveal_around(x, y, mf):
    unchecked_interior = set()
    boundary = set()
    offsets = [(-1, 0), (1, 0), (0, -1), (0, 1)]

    for ox, oy in offsets:
        _x, _y = offset_inside(x, y, ox, oy, mf.shape[0])
        if mf[_x, _y] == 0:
            unchecked_interior.add((_x, _y))
        else:
            boundary.add((_x, _y))
    return unchecked_interior, boundary


def generate_revelation_matrix(x, y, mf):
    reveal = np.zeros(mf.shape)
    if mf[x, y] == -1:
        return None
    elif mf[x, y] > 0:
        reveal[x, y] = 1
        return reveal
    else:
        checked_interior = set()
        checked_interior.add((x, y))
        unchecked_interor, boundary = reveal_around(x, y, mf)
        while len(unchecked_interor) > 0:
            x, y = unchecked_interor.pop()
            int_nbhd, bdry_nbdh = reveal_around(x, y, mf)
            unchecked_interor.update(int_nbhd.difference(checked_interior))
            boundary.update(bdry_nbdh)
            checked_interior.add((x, y))

        interior = list(checked_interior)
        boundary = list(boundary)

        for x, y in interior:
            reveal[x, y] = 1
        for x, y in boundary:
            if mf[x, y] > 0:
                reveal[x, y] = 1
            else:
                print("Mine in boundary")
        return reveal

class Minesweeper:
    def __init__(self, size, mines):
        self.m = mines
        self.s = size
        self.f = mines
        self.lost = False
        self.initialize_field()

    def initialize_field(self):
        minefield = np.zeros((self.s, self.s))
        fog_of_war = np.zeros((self.s, self.s))
        flagged = set()

        mine_positions = sample(range(self.s*self.s), self.m)
        xs = [-1, 0, 1]
        ys = [-1, 0, 1]
        
        for i in mine_positions:
            x, y = linear_to_pair(i, self.s)
            minefield[x, y] = -1
        
        for i in range(self.s*self.s):
            x, y = linear_to_pair(i, self.s)
            if minefield[x, y] == -1:
                continue
            else:
                sub = np.zeros((3, 3))
                for ox in xs:
                    for oy in ys:
                        sub[ox+1, oy+1] = minefield[offset_inside(x, y, ox, oy, self.s)]
                num_mines = np.sum(sub < 0)
                minefield[x, y] = num_mines

        self.mf = minefield
        self.fow = fog_of_war
        self.flagged = flagged

    def flag_at(self, x, y) -> None:
        self.flagged.add((x, y))

    def unflag_at(self, x, y) -> None:
        self.flagged.remove((x, y))

    def reveal_at(self, x, y):
        r = generate_revelation_matrix(x, y, self.mf)
        if r is None:
            self.lost = True
        else:
            self.fow = np.logical_or(self.fow, r)

    def show_minefield(self):
        field = np.abs(self.mf * self.fow)
        return '\n'.join(' '.join(self.colorize(int(i)) for i in row) for row in field)
    
    @classmethod
    def colorize(cls, i: int) -> str:
        return [
            lambda x: f'\033[0m{x}\033[0m',
            lambda x: f'\033[31m{x}\033[0m',
            lambda x: f'\033[32m{x}\033[0m',
            lambda x: f'\033[33m{x}\033[0m',
            lambda x: f'\033[34m{x}\033[0m',
        ][i](i)
    
    def __repr__(self):
        field = np.abs(self.mf * self.fow)
        s = ""

        upper = "  |"
        lower = "  |"
        divider = "---"
        for x in range(self.s):
            u, l = numeral_vertical(x, len(str(self.s)), "b")
            upper = upper + u + "|"
            lower = lower + l + "|"
            divider += "--"
        s = s + upper + "\n" + lower + "\n" + divider + "\n"

        for x in range(self.s):
            s = s + numeral_horizontal(x, len(str(self.s)), "r") + "|"
            for y in range(self.s):
                # if self.mf[x, y] == -1:
                #     s += "m "
                if (x, y) in self.flagged:
                    s += "f "
                elif self.fow[x, y] == 0:
                    s += "x "
                else:
                    s += self.colorize(int(field[x, y]))
                    s += " "
            s = s[:-1]
            s += "|\n"
        return s[:-1] + "\n" + divider
    
    def check_victory(self):
        return (
            len(self.flagged) == self.m                     # all mines are flagged
            and 
            int(np.sum(self.fow)) == (self.s**2 - self.m)   # no stone unturned
        )

def numeral_horizontal(number, space, justification):
    n = str(number)
    if justification == "l":
        s = n + " " * (space - len(n))
    elif justification == "r":
        s = (n[::-1] + " " * (space - len(n)))[::-1]
    elif justification == "c":
        pre = (space - len(n)) // 2
        s = " " * pre + n + " " * (space - len(n) - pre)
    return s

def numeral_vertical(number, space, justification):
    just_map = {"t": "l", "b": "r", "c": "c"}
    n = numeral_horizontal(number, space, just_map[justification])
    return list(n)

def dispatch_choice(ms: Minesweeper, choice) -> None:
    cmd = choice.split(" ")

    if cmd[0] == "f":
        x = int(cmd[1])
        y = int(cmd[2])
        ms.flag_at(x, y)
    elif cmd[0] == "u":
        x = int(cmd[1])
        y = int(cmd[2])
        ms.unflag_at(x, y)
    elif cmd[0] == "r":
        ms.s = int(cmd[1])
        ms.m = int(cmd[2])
        ms.initialize_field()
    elif cmd[0] == "o":
        x = int(cmd[1])
        y = int(cmd[2])
        ms.reveal_at(x, y)
    else:
        print("""Legal actions are:
    "f x y" to place a flag
    "u x y" to remove a flag
    "o x y" to open a tile
    "r s m" o reset the board with size s and m mines
    "q"     to quit
        """)


if __name__ == "__main__":
    # ms = Minesweeper(20, 1)
    is_running = True
    has_won = False

    print("You're playing minesweeper on the torus.")
    size = int(input("What size of board would you like to play? "))
    num_mines = int(input("How many mines would you like? "))
    ms = Minesweeper(size, num_mines)

    while is_running:
        if has_won:
            print("You won!")
            has_won = False
        else:
            print(ms)
        print(f"Flagged {len(ms.flagged)} of {ms.m} mines. Board size is {ms.s}.")
        print(f"Revealed {int(np.sum(ms.fow))} of {ms.s**2} tiles.")
        cmd = input("What would you like to do? ")
        if cmd == "q":
            is_running = False
        else:
            dispatch_choice(ms, cmd)
            has_won = ms.check_victory()