#!/usr/bin/env python3

import math, itertools, sys, os, random
import SAT_lib, my_utils, latin_utils
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--size", type=int, default=8, help="... or order. Default=8")
# Kissat requires 31-bit seed:
parser.add_argument("--seed", type=int, default=my_utils.get_32_bits_from_urandom()&0x7fffffff, help="Seed. Default=random")
parser.add_argument("--normalize", default=False, help="Normalize/reduce (before search). Default=false", action="store_true")
parser.add_argument("--all-solutions", default=False, help="Produce all possible solutions. Default=false", action="store_true")
tmp=" ".join(SAT_lib.SAT_solvers_supported)
#parser.add_argument("--solver", default="libcadical", help="Set SAT solver. Supported: "+tmp)
parser.add_argument("--solver", default="kissat", help="Set SAT solver. Supported: "+tmp)
# verbose for SAT_lib:
parser.add_argument("-v", action="count", help="Increase verbosity. May be -vv, -vvv, -vvvv")
parser.add_argument("--first", type=str, nargs="?", default=None, help="Set first square (in short form)")
# it produces proof that there is no 2-MOLS or if no more 2-MOLS exist
parser.add_argument("--proof", help="Generate proof. Default=false", action="store_true")

args=parser.parse_args()

# https://stackoverflow.com/questions/4042452/display-help-message-with-python-argparse-when-script-is-called-without-any-argu
parser.print_help()
print ("")

_seed=args.seed
print ("Setting seed", _seed)
random.seed(_seed)

if args.first!=None:
    print ("Setting first square")
    sq=args.first
    SIZE=int(math.sqrt(len(sq)))
else:
    SIZE=args.size
print ("Setting size", SIZE)

# both squares? second only?
if args.normalize:
    print ("Normalize mode")
else:
    print ("Non-normalize mode")

if args.proof:
    print ("Proof will be produced")
    proof=True
else:
    proof=False

if args.v!=None:
    verbosity=args.v
else:
    verbosity=0

SAT_solver=args.solver
print ("Setting solver", SAT_solver)

s=SAT_lib.SAT_lib(seed=_seed, SAT_solver=SAT_solver, proof=proof, threads=10)

a=[[s.alloc_BV(SIZE) for _ in range(SIZE)] for _ in range(SIZE)]
b=[[s.alloc_BV(SIZE) for _ in range(SIZE)] for _ in range(SIZE)]

latin_utils.latin_add_constraints(SIZE, s, a)
latin_utils.latin_add_constraints(SIZE, s, b)

latin_utils.make_mutually_orthogonal(SIZE, s, a, b)

if args.first==None and args.normalize:
    latin_utils.fix_first_row_increasing(s, a, SIZE)
    latin_utils.fix_first_col_increasing(s, a, SIZE)
    latin_utils.fix_first_col_increasing(s, b, SIZE)

# if first is set, normalize only mate:
if args.first!=None and args.normalize:
    latin_utils.fix_first_col_increasing(s, b, SIZE)

if args.first!=None:
    print ("Setting first square")
    sq=args.first
    sq=latin_utils.base36_str_to_list_of_lists(sq)
    if args.normalize:
        sq=latin_utils.normalize(sq)
        print ("Normalizing first square, now it's: "+latin_utils.list_of_lists_of_int_to_base36_str(sq))
    latin_utils.fix_square_to_hardcoded(s, a, sq, SIZE)

sol=0

if s.solve():
    while True:
        sol=sol+1
        print ("*** Solution %d" % sol)

        print ("First:")
        latin_utils.print_LS_from_SAT_vars(s, a)
        if SIZE<=36:
            print ("Short form:", latin_utils.list_of_lists_of_int_to_base36_str(latin_utils.get_square(s, a)))
        print ("")
        print ("Mate:")
        latin_utils.print_LS_from_SAT_vars(s, b)
        if SIZE<=36:
            print ("Short form:", latin_utils.list_of_lists_of_int_to_base36_str(latin_utils.get_square(s, b)))
        print ("")
        print ("Concatenated:")
        latin_utils.print_MOLS2_from_SAT_vars(s, a, b)
        print ("")
        sys.stdout.flush()

        if args.all_solutions==False:
            break
        if s.fetch_next_solution()==False:
            break
else:
    print ("UNSAT")

print ("Solutions: %d" % sol)