import random
from IPython.core.display import SVG
import pyomo.environ as pyo
from pysat.solvers import Solver
from pysat.formula import CNF 
import py_svg_combinatorics as psc
from ipywidgets import widgets, HBox
from collections import Counter
from pprint import pprint
from random import randint
import numpy as np
from IPython.display import IFrame
from IPython.display import display
import IPython
from copy import copy
import os
from pyomo_helpers import * 
import pandas
from pathlib import Path
from pyomo_helpers import pretty_solution
nbname = ''
try:
    nbname = __vsc_ipynb_file__
except:
    if 'COCALC_JUPYTER_FILENAME' in os.environ:
        nbname = os.environ['COCALC_JUPYTER_FILENAME']    
title_ = Path(nbname).stem.replace('-', '_').title()    
IFrame(f'https://discopal.ispras.ru/index.php?title=Hardprob/{title_}&useskin=cleanmonobook', width=1280, height=300)

d = 10
n = 5 # Number of vectors in X
M = 4

vectors = np.random.randint(0, 2, size=(d, n))

def visualize(vectors):
    df = pd.DataFrame(vectors, columns=["x" + str(i) for i in range(vectors.shape[1])])
    s = df.style.set_table_styles( [{'selector': 'td',
        'props': [
            ('border-left', '2px solid black !important'),
            ('border-right','2px solid black !important')
        ]}])
    return s
visualize(vectors)

def модель(вектора, M):
    m = pyo.ConcreteModel()
    m.множества = range(M)
    m.В = range(вектора.shape[1])
    m.D = range(вектора.shape[0])
    m.вектора = вектора
    m.X = pyo.Var(m.В, m.множества, domain=pyo.Binary)
    m.является_покрытием = pyo.Var(m.множества, domain=pyo.Binary) 
    
    @m.Objective(sense=pyo.maximize)
    def количество_покрытий(m, в):
        return sum(m.является_покрытием[...]) # важно указывать слайс при суммировании!
    
    @m.Constraint(m.В)
    def ровно_в_одном_множестве(m, в):
        return sum(m.X[в, ...]) == 1
    
    @m.Constraint(m.множества, m.D)
    def проверка_является_покрытием(m, мн, d):
        #Если хоть кто-то в мн 0 - то мн не покрытие
        return (m.вектора @ m.X)[d, мн] >= m.является_покрытием[мн]

    # @m.Constraint()
    # def wtf(m):
    #     return m.X[1,0]==1

    return m

m = модель(vectors, M)
m.pprint()

solver = pyo.SolverFactory('cbc')
# solver = pyo.SolverFactory('scip')
solver.solve(m).write()

def solve_d_vectors(vectors, M, solver_name='cbc'):
    m = модель(vectors, M)
    solver = pyo.SolverFactory(solver_name)
    solver.solve(m)
    cover = np.array([pyo.value(m.является_покрытием[i]) for i in range(M)])
    solution = np.array([[int(pyo.value(m.X[i, j])) for j in range(M)] for i in range(vectors.shape[1])])
    solution = solution[..., cover > 0]
    answer = m.количество_покрытий()
    return answer, solution

def visualize_solution(solution):
    df = pd.DataFrame(solution, columns=["s" + str(i) for i in range(solution.shape[1])], index=["x" + str(i)  for i in range(solution.shape[0])])
    return df.style.background_gradient()

visualize_solution(solve_d_vectors(vectors, M)[1])

cnf_vars = 5
cnf_clauses = 20

cnf3 = CNF(from_clauses=psc.rand3cnf(cnf_clauses, cnf_vars))
clauses = cnf3.clauses

def solve_sat(cnf_vars, cnf_clauses, clauses):
    vectors = np.zeros((cnf_clauses + cnf_vars, 2*cnf_vars + 1))
    for i in range(cnf_clauses):
        for j in clauses[i]:
            vectors[i][j + cnf_vars] = 1
        vectors[i][cnf_vars] = 1 # Contradict

    for i in range(1, cnf_vars+1):
        vectors[cnf_clauses + i - 1][cnf_vars - i] = 1
        vectors[cnf_clauses + i - 1][cnf_vars + i] = 1 
    sol = solve_d_vectors(vectors, 2)
    return sol[0] == 2

solve_sat(cnf_vars, cnf_clauses, clauses)

cnf_vars = 5
cnf_clauses = 1

cnf3 = CNF(from_clauses=psc.rand3cnf(cnf_clauses, cnf_vars))
clauses = cnf3.clauses

clauses

solve_sat(cnf_vars, cnf_clauses, clauses)

def test_sat_iter(N, cnf_vars, cnf_clauses):
    mistakes = 0
    for _ in range(N):
        cnf3 = CNF(from_clauses=psc.rand3cnf(cnf_clauses, cnf_vars))
        clauses = cnf3.clauses
        sol_this = solve_sat(cnf_vars, cnf_clauses, clauses)
        py3sat_test_solver = Solver(bootstrap_with=cnf3)
        sol_other = py3sat_test_solver.solve()
        if sol_this != sol_other:
            mistakes += 1
    return mistakes

def stress_test():
    for i in range(100):
        if test_sat_iter(20, 4, 20) > 5:
            print("Failed")
            return
        if i % 10 == 0:
            print("Passed ", i)

stress_test()