#*****************************************************************************
#       Copyright (C) 2007 Carlo Hamalainen <carlo.hamalainen@gmail.com>, 
#
#  Distributed under the terms of the GNU General Public License (GPL)
#
#    This code is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#    General Public License for more details.
#
#  The full text of the GPL is available at:
#
#                  http://www.gnu.org/licenses/
#*****************************************************************************

import copy

import sage.rings.integer_mod_ring

cdef extern from "stdlib.h":
    ctypedef unsigned long size_t
    void *malloc(size_t size)
    void free(void *mem)

cdef class latin_square:
    """
    This class represents a latin square as a 1D
    array and provides the usual access methods. 
    """

    cdef int *square

    # self.n is maintained as the order of the latin square.
    cdef readonly int n

    def __new__(self, n = 0, clear_square = True):
        cdef int i

        self.n = n

        if n > 0:
            self.square = <int*> malloc(sizeof(int)*n*n)

            if clear_square: 
                for i from 0 <= i < n*n:
                    self.square[i] = -1

        else:
            self.square = NULL

    def __dealloc__(self):
        free(self.square)

    def __getitem__(self, rc):
        """
        If L is a latin_square then this method
        allows us to evaluate L[r, c].
        """

        assert self.n > 0

        r = rc[0]
        c = rc[1]

        return self.square[self.n*r + c];

    def __setitem__(self, rc, val):
        """
        If L is a latin_square then this method
        allows us to set L[r, c].
        """
        assert self.n > 0

        r = rc[0]
        c = rc[1]

        self.square[self.n*r + c] = val

    def __deepcopy__(self, memo = {}):
        cdef int i, j
        cdef object L_deepcopy

        L_deepcopy = latin_square(self.n)

        for i in range(self.n):
            for j in range(self.n):
                L_deepcopy[i, j] = self[i, j]

        return L_deepcopy

    def __str__(self):
        s = ""

        for i in range(self.n):
            for j in range(self.n):
                if self[i, j] < 0:
                    s += "- "
                else:
                    s += str(self[i, j]) + " "
            s += "\n"
        return s

    def size(self):
        raise NotImplemented

    def nrows(self):
        return self.n

    def ncols(self):
        return self.n

    def permissable_values(self, int r, int c):
        cdef int i, j, e

        assert self[r, c] < 0

        vals = {}
        for e from 0 <= e < self.n:
            vals[e] = 1

        for i from 0 <= i < self.n:
            if self[i, c] >= 0:
                del vals[ self[i, c] ]

        for j from 0 <= j < self.n:
            if self[r, j] >= 0:
                try:
                    del vals[ self[r, j] ]
                except KeyError:
                    # We may have already removed a symbol
                    # in the previous for-loop.
                    pass

        return vals.keys()

cdef top_left_empty_cell(latin_square L):
    cdef int i, j

    for i from 0 <= i < L.n:
        for j from 0 <= j < L.n:
            if L[i, j] < 0: return (i, j)

    return None

cdef random_empty_cell_pyrex(latin_square P):
    cdef object R

    cells = {}

    for r in range(P.nrows()):
        for c in range(P.ncols()):
            if P[r, c] < 0:
                cells[ (r,c) ] = True

    cells = cells.keys()

    if len(cells) == 0: return None

    R = sage.rings.integer_mod_ring.IntegerModRing()

    #rc = cells[ R.random_element(len(cells)) ]
    rc = cells[ R.random_element( len(cells) ) ]

    return [rc[0], rc[1]]


cdef dfs_search_internal(latin_square L, random_search, int nr_to_find, completions):
    """
    Basic depth first search algorithm for finding
    completion(s) of a partial latin square.

    INPUT:
        L --                        Partial latin square
        random_search --            If this evalutes to True we choose
                                    an unfilled cell uniformly at
                                    random.
        nr_to_find --               Maximum number of completions to find.
        completions --              Completed latin squares are added to
                                    this dictionary.

    EXAMPLE:
        See dfs_pyrex() below for an example of how to use the
        dfs_search_internal function.
    """

    cdef int i, j
    cdef object L_new

    if random_search:
        ij = random_empty_cell_pyrex(L)
    else:
        ij = top_left_empty_cell(L)

    if ij == None:
        L_new = latin_square(L.n, clear_square = False)
        L_new = copy.deepcopy(L)

        completions[L_new] = True

        return

    else:
        i = ij[0]
        j = ij[1]

    vals = L.permissable_values(i, j)

    for e in vals:
        L_new = latin_square(L.n, clear_square = False)
        L_new = copy.deepcopy(L)

        # We are trying to fill in the [i, j] cell of
        # the latin square L.
        assert L[i, j] < 0

        L_new[i, j] = e
        dfs_search_internal(L_new, random_search, nr_to_find, completions)

        if nr_to_find > 0 and len(completions) >= nr_to_find:
            return

def dfs_pyrex(L_in, random_search, nr_to_find):
    """
    Depth first search in Pyrex for finding 
    completions of a partial latin square.

    INPUT:
        L_in -- matrix,             partial latin square of type matrix
                                    with (-1) indicating empty cells.
        random_search --            If this evalutes to True we choose
                                    an unfilled cell uniformly at
                                    random.
        nr_to_find -- int,          Maximum number of completions to find.

    NOTES:
        There are much better ways of finding partial latin square
        completions, e.g. convert to a 0-1 matrix exact cover representation
        and use Knuth's Dancing Links algorithm.
    """

    cdef object L
    cdef int r, c, n

    n = L_in.nrows()
 
    completions = {}

    L = latin_square(n)

    for r from 0 <= r < n:
        for c from 0 <= c < n:
            L[r, c] = L_in[r, c]

    dfs_search_internal(L, random_search, nr_to_find, completions)

    assert len(completions) <= nr_to_find

    return completions.keys()