def paint_numbers(numbers, count):
    if sum(numbers) + len(numbers) - 1 > count: #❶
        return
    elif len(numbers) == 1: #❷
        for i in range(0, count - numbers[0] + 1):
            yield [-1] * i + [1] * numbers[0] + [-1] * (count - i - numbers[0])
    else:
        for res in paint_numbers(numbers[1:], count - numbers[0] - 1): #❸
            yield [1] * numbers[0] + [-1] + res
        for res in paint_numbers(numbers, count - 1): #❹
            yield [-1] + res

for line in paint_numbers([2, 1, 5], 12):
    print([int(x > 0) for x in line])

from sympy import symbols, Equivalent
from sympy.logic.boolalg import to_cnf

X, A, B, C = symbols("X A B C")
expr = Equivalent(X, ~A & B & C)
to_cnf(expr)

def dnf_to_cnf(dnf, new_vars=None):
    if new_vars is None:
        start = max(max(map(abs, term)) for term in dnf) + 1
        new_vars = iter(range(start, start + len(dnf)))
    zlist = []
    cnf = []
    for term in dnf:
        z = next(new_vars)
        zlist.append(z)
        cnf.append([z] + [-v for v in term])
        for v in term:
            cnf.append([-z, v])
    cnf.append(zlist)
    return cnf

from pprint import pprint
from pysat.solvers import Solver

dnf = [[-1, 2, 3], [1, -2, 3], [1, 2, -3]]
cnf = dnf_to_cnf(dnf)
print("CNF")
pprint(cnf)
print("solutions")
sat = Solver()
sat.append_formula(cnf)
for sol in sat.enum_models():
    print([int(x > 0) for x in sol])

n = 12
solution = list(paint_numbers([2, 1, 5], n))
variables = list(range(1, n+1))
dnf = []
for term in solution:
    dnf.append([v * t for v, t in zip(variables, term)])
sat = Solver()
sat.append_formula(dnf_to_cnf(dnf))
solution_cnf = [[x // abs(x) for x in res[:n]] for res in sat.enum_models()]
sorted(solution) == sorted(solution_cnf)

import numpy as np
import itertools

class NonogramSimpleSolver:
    def __init__(self, rows, cols, width=None, height=None):
        if width is None:
            width = len(cols)
        if height is None:
            height = len(rows)
        self.rows = rows
        self.cols = cols
        self.width = width
        self.height = height
        self.var_count = self.width * self.height
        self.vars = np.arange(1, self.var_count + 1).reshape(-1, self.width) #❶
        self.new_vars = itertools.count(self.var_count + 2) #❷
        
        self.cnf = []
        
        for i, row in enumerate(self.rows):
            self.add_constraints(row, self.width, self.vars[i, :].tolist())
            
        for j, col in enumerate(self.cols):
            self.add_constraints(col, self.height, self.vars[:, j].tolist())
            
    def add_constraints(self, numbers, count, variables): #❸
        dnf = []
        for pattern in paint_numbers(numbers, count):
            dnf.append([v * p for v, p in zip(variables, pattern)])
        self.cnf.extend(dnf_to_cnf(dnf, new_vars = self.new_vars))
        
    def solve(self): #❹
        sat = Solver()
        sat.append_formula(self.cnf)
        sat.solve()
        model = sat.get_model()
        self.solution = np.array(model[:self.var_count]).reshape(self.height, -1)
        self.solution = np.sign(self.solution)
        self.solution[self.solution == -1] = 0
    
    def plot(self, font_size=9): #❺
        from helper.plot import plot_nonogram_board
        return plot_nonogram_board({'rows':self.rows, 'cols':self.cols}, self.solution, font_size=font_size)

rows = [
    [3],
    [2,1],
    [2],
    [7],
    [9],
    [2, 3, 2],
    [4, 4],
    [9],
    [2, 3, 2],
    [3, 3],
    [7],
    [5]
]

cols = [
    [6],
    [8],
    [2, 2, 3],
    [8, 2],
    [6, 2, 2],
    [1, 6, 2],
    [2, 2, 2, 3],
    [8],
    [6]
]
pbn = NonogramSimpleSolver(rows, cols)
pbn.solve()
pbn.plot();

import json

def solve_and_paint(fn, solver, font_size=9):
    if isinstance(fn, dict):
        data = fn
    else:
        with open(fn) as f:
            data = json.load(f)
            
    pbn = solver(data["rows"], data["cols"])
    pbn.solve()
    pbn.plot(font_size=font_size)

%time solve_and_paint("data/nonogram02.json", solver=NonogramSimpleSolver)

len(list(paint_numbers([1, 2, 3, 2, 1], 30)))

def start_pos(numbers, count):
    s = 0
    for i, n in enumerate(numbers):
        yield s
        s += n + 1

def end_pos(numbers, count):
    width = sum(numbers) + len(numbers) - 1
    start = count - width
    for pos in start_pos(numbers, count):
        yield start + pos
        
def create_position_variables(numbers, count, new_bool_variables):
    pos_variables = []
    for i, (a, b) in enumerate(zip(start_pos(numbers, count), end_pos(numbers, count))):
        pos_variables.append({j:next(new_bool_variables) for j in range(a, b + 1)})
    return pos_variables

numbers = [1, 2, 3]
count = 10
pos_variables = create_position_variables(numbers, count, itertools.count(1))
pos_variables

for line in paint_numbers([1, 2, 3], 10):
    print([int(x > 0) for x in line])

def only_one_true(variables):
    yield list(variables)
    for i, j in itertools.combinations(variables, 2):
        yield [-i, -j]
        
def pairwise(iterable):
    a, b = itertools.tee(iterable)
    next(b, None)
    return zip(a, b)        

def position_variables_constraint(pos_variables, numbers):
    cnf = []
    # 各数字の位置変数のうち1つだけがTrue
    for i, pos_var in enumerate(pos_variables):
        cnf.extend(only_one_true(pos_var.values()))

    # 連続する数字の位置関係制約
    for i, (pos_var1, pos_var2) in enumerate(pairwise(pos_variables)):
        width = numbers[i] + 1
        for pos1, pos2 in itertools.product(pos_var1, pos_var2):
            if pos2 < pos1 + width:
                cnf.append([-pos_var1[pos1], -pos_var2[pos2]])
    return cnf

cnf = position_variables_constraint(pos_variables, numbers)
cnf

sat = Solver()
sat.append_formula(cnf)

cell_map = {}
for n, pos_var in zip(numbers, pos_variables):
    for key, val in pos_var.items():
        cell_map[val] = (key, n)

for sol in sat.enum_models():
    true_vars = [loc for loc, v in enumerate(sol, 1) if v > 0]
    cells = [0] * 10
    for v in true_vars:
        s, n = cell_map[v]
        for j in range(s, s+n):
            cells[j] = 1
    print(sol, true_vars, cells)

def cell_variables_constraint(pos_variables, cell_variables, numbers, count):
    cnf = []
    # cell[j] == 1 for all pos[i] <= j < pos[i] + num[i]
    for i, pos_var in enumerate(pos_variables):
        for pos, var in pos_var.items():
            for j in range(numbers[i]):
                cnf.append([-var, cell_variables[pos + j]])

    # cell[j] == 0, for all j < pos[0]
    for pos, var in pos_variables[0].items():
        for j in range(0, pos):
            cnf.append([-var, -cell_variables[j]])

    # cell[j] == 0, for all j >= pos[-1] + num[-1]
    for pos, var in pos_variables[-1].items():
        for j in range(pos + numbers[-1], count):
            cnf.append([-var, -cell_variables[j]])

    # cell[j] == 0 for all pos[i] + num[i] <= j < pos[i+1]
    for i, (pos_var1, pos_var2) in enumerate(pairwise(pos_variables)):
        for pos1, pos2 in itertools.product(pos_var1, pos_var2):
            if pos2 > pos1 + numbers[i]:
                for j in range(pos1 + numbers[i], pos2):
                    cnf.append([-pos_var1[pos1], -pos_var2[pos2], -cell_variables[j]])
    return cnf

numbers = [1, 2, 3]
count = 10
cell_variables = list(range(1, count + 1))
pos_variables = create_position_variables(numbers, count, itertools.count(count + 1))
cnf = position_variables_constraint(pos_variables, numbers)
cnf += cell_variables_constraint(pos_variables, cell_variables, numbers, count)

sat = Solver()
sat.append_formula(cnf)
for sol in sat.enum_models():
    print([int(x > 0) for x in sol[:count]])

import numpy as np

class NonogramFastSolver(NonogramSimpleSolver):            
    def add_constraints(self, numbers, count, cell_variables):
        pos_variables = create_position_variables(numbers, count, self.new_vars)
        self.cnf.extend(position_variables_constraint(pos_variables, numbers))
        self.cnf.extend(cell_variables_constraint(pos_variables, cell_variables, numbers, count))

%time solve_and_paint("data/nonogram02.json", solver=NonogramFastSolver)

%time solve_and_paint("data/nonogram03.json", solver=NonogramFastSolver, font_size=6)