"""
Random latin square generator using the Method of Circuts.
Model implementation in Python with symbol indexing.
Larry Trammell -- May 2020

This supports the document "Direct Construction of Random Latin 
Squares by the Method of Circuits" and is publicly distributed
under the Creative Commons CC-BY-4.0 license.
"""

# Define row-column dimensions of square to be constructed
Nrc = 10

# Python system imports
import random as rnpy            # core random number generator
import numpy as np               # numeric data types
from copy import deepcopy        # data replicator for initialization   

"""
Main generator sequence
"""
        
print("Generating a Latin Square of size ", Nrc,"x",Nrc)
empty = np.short(-1)

# Initialize an empty square
A = np.zeros((Nrc,Nrc), dtype=np.short)

# Initialize feasible symbols sets for rows and cols
row = [ ]
for j in range(Nrc) :
    row.append( np.short(j+1) )
rowavail = [ ]
for i in range(Nrc) :
    rowavail.append( set(row) )
colavail = deepcopy(rowavail)

# Initialize lookup indexing tables by row and column
rowx = np.ndarray((Nrc,Nrc+1), dtype=np.short)
colx = np.ndarray((Nrc,Nrc+1), dtype=np.short)
rowx.fill(np.short(empty))
colx.fill(np.short(empty))

# Assigning symbols to layout cells in row-column lexical order. 
for i in range(Nrc) :
    for j in range(Nrc):

        # Check for row-column feasible symbols at this new location
        avail = rowavail[i] & colavail[j]

        if len(avail) > 0 :
            # FEASIBLE. Assign a randomly-selected feasible symbol.
            sel=rnpy.sample(avail,1)
            symbol1 = sel[0]
            A[i,j] = symbol1
            # Remove symbol from availability sets
            rowavail[i].remove(symbol1)
            colavail[j].remove(symbol1)
            # Index the symbol for row and column lookup
            rowx[i,symbol1] = j
            colx[j,symbol1] = i

        else :
            # CONFLICT. Insert a randomly-selected row-feasible SYMBOL1 
            sel=rnpy.sample(rowavail[i],1)
            symbol1 = sel[0]
            A[i,j] = symbol1
            rowavail[i].remove(symbol1)
            if symbol1 in colavail[j] :
                colavail[j].remove(symbol1)

            # Index SYMBOL1 for starting row lookup
            rowx[i,symbol1] = j

            # Randomly select a column-feasible replacement SYMBOL2
            sel=rnpy.sample(colavail[j],1)
            symbol2 = sel[0]
            
            # Begin conflict resolution at the newly-assigned cell (i,j).          
            iviol = i
            jviol = j

            while(True) :
                    
                # Locate row of the conflicting entry in column jviol
                iprev = iviol
                iviol = colx[jviol,symbol1]
                colx[jviol,symbol1] = iprev
                colx[jviol,symbol2] = empty          
                # If no conflict, resolution is complete
                if  iviol == empty :
                    break

                # Replace SYMBOL1 with SYMBOL2 at the current location
                A[iviol,jviol] = symbol2
                if symbol2 in rowavail[iviol] :
                    rowavail[iviol].remove(symbol2)
                rowavail[iviol].add(symbol1)
                colavail[jviol].remove(symbol2)
                colx[jviol,symbol2] = iviol

                # Locate column of the violating symbol in row iviol
                jprev = jviol
                jviol = rowx[iviol,symbol2]
                rowx[iviol,symbol2] = jprev
                rowx[iviol,symbol1] = empty          
                # If no conflict, resolution is complete
                if  jviol == empty :
                    break

                # Replace SYMBOL2 with SYMBOL1 at the current location
                A[iviol,jviol] = symbol1
                if symbol1 in colavail[jviol] :
                    colavail[jviol].remove(symbol1)
                colavail[jviol].add(symbol2)
                rowavail[iviol].remove(symbol1)
                rowx[iviol,symbol1] = jviol
            pass
        pass
    pass
       
print("\nCompleted square ")
print(A)

#print("\nVerify row and column sums")
#rowsum = [0]*Nrc
#colsum = [0]*Nrc
#for i in range(Nrc) :
#    for j in range(Nrc) :
#        rowsum[i] = rowsum[i]+A[i,j]
#        colsum[j] = colsum[j]+A[i,j]
#print("rowsum :", rowsum)
#print("colsum :", colsum)