import json
import sys
import os
import multiprocessing as mp
from contextlib import redirect_stdout
import torch
import numpy as np

from solver.Crossword import Crossword
from solver.BPSolver import BPSolver
from solver.Utils import print_grid, convert_puz

def run_examples(batch):
    sys.path.append('./solver')
    for args in batch:
        puzzle_file, logdir = args
        if puzzle_file[-4:] == '.puz':
            try:
                puzzle = convert_puz(puzzle_file)
            except:
                print('failed on', puzzle_file)
                continue
        else:
            with open(puzzle_file, 'r') as infile:
                puzzle = json.load(infile)
        with open(os.path.join(logdir, os.path.split(puzzle_file)[1] + '.txt'), 'w') as wf:
            curr_process = 0
            try:
                print('running solver on file', puzzle_file)
                crossword = Crossword(puzzle)
                print_grid(crossword.letter_grid)
                solver = BPSolver(crossword, max_candidates=500000, process_id=curr_process) # use 500000 for best results
                solution = solver.solve(num_iters=10, iterative_improvement_steps=5)
                print_grid(solution)
                solver.evaluate(solution)
            except:
                print('PUZZLE FAILED', puzzle_file)

logdir ='logs/testing/'
os.makedirs(logdir, exist_ok=True)

run_examples([('solver/Mar2321.puz', logdir)])
