#!/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)")
parser.add_argument("--tpl-second", type=str, nargs="?", default=None, help="Set template for second square (in short form, dots are for any symbol in cell)")
# 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")
parser.add_argument("--time", type=int, default=-1, help="Timeout in seconds. Default=none")

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)

timeout=-1
if args.time!=-1:
    timeout=args.time
    print ("Setting timeout:", timeout)

if args.tpl_second!=None:
    print ("Setting template for second square")
    tpl2=args.tpl_second
else:
    tpl2=None

# 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, timeout=timeout, verbose=0)

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)]
c=[[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.latin_add_constraints(SIZE, s, c)

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

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)
    latin_utils.fix_first_col_increasing(s, c, SIZE)

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

if args.first!=None:
    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)
    print ("transv_cnt=%05d" % latin_utils.count_transversals(sq))
print (f"PID={os.getpid()}")
sys.stdout.flush()

if tpl2!=None:
    for _r in range(SIZE):
        for _c in range(SIZE):
            if len(tpl2)-1>=(_r*SIZE+_c):
                x=tpl2[_r*SIZE+_c]
                if x!='.':
                    symb=ord(x)-ord('0')
                    tmp=SAT_lib.n_to_BV(1 << symb, SIZE)
                    #print (_r, _c, tmp)
                    s.fix_BV(b[_r][_c], tmp)
                    print (f"Fixing second square. {_r=} {_c=} {symb=}")

sol=0

try:
    res=s.solve()
except SAT_lib.SATSolverTimeout:
    print ("TIMEOUT")
    exit(0)
if res==False:
    print ("UNSAT")
if res:
    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 1:")
        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 ("Mate 2:")
        latin_utils.print_LS_from_SAT_vars(s, c)
        if SIZE<=36:
            print ("Short form:", latin_utils.list_of_lists_of_int_to_base36_str(latin_utils.get_square(s, c)))
        print ("")
        print ("Concatenated:")
        latin_utils.print_MOLS3_from_SAT_vars(s, a, b, c)
        print ("")
        sys.stdout.flush()

        if args.all_solutions==False:
            break
        if s.fetch_next_solution()==False:
            break

print ("Solutions: %d" % sol)