# TODO:
# Refactor to make it more Pythonic
# Use functions directly instead of utils
# Reimplement util functions
# Use assign_value to visualize
# Optimize naked_twins to get rid of duplicated cases
# Implement other strategies
# Implement naked 

# import utils

assignments = []

rows = 'ABCDEFGHI'
cols = '123456789'

def cross(a, b):
    return [s+t for s in a for t in b]

boxes = cross(rows, cols)

row_units = [cross(r, cols) for r in rows]
column_units = [cross(rows, c) for c in cols]
square_units = [cross(rs, cs) for rs in ('ABC','DEF','GHI') for cs in ('123','456','789')]

diagonal_units = [list(map(lambda x,y: x+y, rows, all_cols)) for all_cols in [cols, cols[::-1]]]

unitlist = row_units + column_units + square_units + diagonal_units
units = dict((s, [u for u in unitlist if s in u]) for s in boxes)
peers = dict((s, set(sum(units[s],[]))-set([s])) for s in boxes)

def assign_value(values, box, value):
    """
    Please use this function to update your values dictionary!
    Assigns a value to a given box. If it updates the board record it.
    """
    values[box] = value
    if len(value) == 1:
        assignments.append(values.copy())
    return values

def display(values):
    """
    Display the values as a 2-D grid.
    Input: The sudoku in dictionary form
    Output: None
    """
    width = 1+max(len(values[s]) for s in boxes)
    line = '+'.join(['-'*(width*3)]*3)
    for r in rows:
        print(''.join(values[r+c].center(width)+('|' if c in '36' else '')
                      for c in cols))
        if r in 'CF': print(line)
    print

def grid_values(grid):
    """
    Convert grid into a dict of {square: char} with '123456789' for empties.
    Input: A grid in string form.
    Output: A grid in dictionary form
            Keys: The boxes, e.g., 'A1'
            Values: The value in each box, e.g., '8'. If the box has no value, then the value will be '123456789'.
    """
    return dict(zip(boxes,map(lambda x: ('123456789' if x=='.' else x), [char for char in grid])))
    
def eliminate(values):
    """
    Go through all the boxes, and whenever there is a box with a value, eliminate this value from the values of all its peers.
    Input: A sudoku in dictionary form.
    Output: The resulting sudoku in dictionary form.
    """

    # For each blank box (value='123456789'), scan for its peers with predefined values and eliminate those values from it
    new_values = values.copy()
    for box in new_values:
        if len(values[box]) != 1:
            for peer in peers[box]:
                if (len(values[peer])) == 1:
                    assign_value(new_values, box, new_values[box].replace(values[peer],''))
    return new_values

    # Alternate solution by Udacity: iterate through solved values then remove that value from its peers (more straightforward)
    
    # solved_values = [box for box in values.keys() if len(values[box]) == 1]
    # for box in solved_values:
    #     digit = values[box]
    #     for peer in peers[box]:
    #         # assign_value(values,peer,values[peer].replace(digit,''))
    #         values[peer] = values[peer].replace(digit,'')
    # return values

def only_choice(values):
    """
    Go through all the units, and whenever there is a unit with a value that only fits in one box, assign the value to this box.
    Input: A sudoku in dictionary form.
    Output: The resulting sudoku in dictionary form.
    """
    new_values = values.copy()
    for unit in unitlist:
        for digit in '123456789':
            dplaces = [box for box in unit if digit in values[box]]
            if len(dplaces) == 1:
                assign_value(new_values, dplaces[0], digit)
                # new_values[dplaces[0]] = digit
    return new_values

def naked_twins(values):
    """Eliminate values using the naked twins strategy.
    Args:
        values(dict): a dictionary of the form {'box_name': '123456789', ...}

    Returns:
        the values dictionary with the naked twins eliminated from peers.
    """

    # For each box, scan for naked twins in its peers
    # Once found, find affected peers, which is intersection between the peers of current box and the peers of the twin box
    # Remove all naked twin values from affected peers
    
    for box in values:
        if len(values[box])==2:
            for peer in peers[box]:
                if values[peer]==values[box]: # found a naked twin
                    for ap in peers[box] & peers[peer]:
                        assign_value(values, ap, values[ap].replace(values[box][0],'').replace(values[box][1],''))
                        # values[ap] = values[ap].replace(values[box][0],'').replace(values[box][1],'')
    return values

def reduce_puzzle(values):
    """
    Iterate eliminate() and only_choice(). If at some point, there is a box with no available values, return False.
    If the sudoku is solved, return the sudoku.
    If after an iteration of both functions, the sudoku remains the same, return the sudoku.
    Input: A sudoku in dictionary form.
    Output: The resulting sudoku in dictionary form.
    """
    solved_values = [box for box in values.keys() if len(values[box]) == 1]
    stalled = False
    while not stalled:
        solved_values_before = len([box for box in values.keys() if len(values[box]) == 1])
        values = eliminate(values)
        values = only_choice(values)
        values = naked_twins(values)
        solved_values_after = len([box for box in values.keys() if len(values[box]) == 1])
        stalled = solved_values_before == solved_values_after
        if len([box for box in values.keys() if len(values[box]) == 0]):
            return False
    return values

def search(values):
    # Using depth-first search and propagation, create a search tree and solve the sudoku.
    # First, reduce the puzzle using the previous function
    
    # Choose one of the unfilled squares with the fewest possibilities
    
    # Now use recursion to solve each one of the resulting sudokus, and if one returns a value (not False), return that answer!

    # If you're stuck, see the solution.py tab!
    
    values = reduce_puzzle(values)
    if values is False:
        return False ## Failed earlier
    if all(len(values[s]) == 1 for s in boxes): 
        return values ## Solved!
    n,start = min((len(values[s]), s) for s in boxes if len(values[s]) > 1)
    
    for digit in values[start]:
        new_values = values.copy()
        assign_value(new_values, start, digit)
        # new_values[start] = digit
        attempt = search(new_values)
        if attempt:
            return attempt

def solve(grid):
    """
    Find the solution to a Sudoku grid.
    Args:
        grid(string): a string representing a sudoku grid.
            Example: '2.............62....1....7...6..8...3...9...7...6..4...4....8....52.............3'
    Returns:
        The dictionary representation of the final sudoku grid. False if no solution exists.
    """
    return search(grid_values(grid))

if __name__ == '__main__':
    diag_sudoku_grid = '2.............62....1....7...6..8...3...9...7...6..4...4....8....52.............3'
    display(solve(diag_sudoku_grid))

    try:
        from visualize import visualize_assignments
        visualize_assignments(assignments)

    except SystemExit:
        pass
    except Exception as e: 
        print(e)
    # except:
    #     print('We could not visualize your board due to a pygame issue. Not a problem! It is not a requirement.')
