from __future__ import annotations
import typing
import random
from IPython.core.display import SVG
import pyomo.environ as pyo
import pyomo.util.infeasible
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
import IPython
from copy import copy
import os
from pathlib import Path
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)

def gen_random_inputs(k: int, n_vars: int, n_terms: int) -> np.ndarray['n_terms,k']:
    """
    Returns C - the matrix of the CNF ter

    A var is an integer in [1, n_vars]
    A literal is a var, -var or 0, which means False for disjunctions (i.e. no effect)
    
    The set of all vars is implicitly defined as np.arange(n_vars) + 1
    """
    
    C: np.ndarray['n_terms,k'] = np.random.randint(n_vars + 1, size=(n_terms, k))
    C *= np.random.choice([-1, 1], size=(n_terms, 3))
    
    return C

k = 3
n_vars = 10
n_terms = 50
C = gen_random_inputs(k, n_vars, n_terms)

k, n_vars, n_terms, C

def get_model(k: int, n_vars: int, n_terms: int, C: np.ndarray['n_terms,k']) -> pyo.ConcreteModel:
    m = pyo.ConcreteModel(name="maximum k-satisfiability")

    m.n_vars = n_vars
    m.n_terms = n_terms
    m.k = k
    
    """
    C holds the matrix of the CNF terms.
    Each row is a term, each column is a literal.
    A literal is a var, -var or 0, which means False for disjunctions (i.e. no effect)
    A var is an integer in [1, n_vars]
    """
    assert C.shape == (n_terms, k)
    m.C = C

    """
    var_values holds the truth values of the variables.
    """
    m.I = range(m.n_vars)
    m.J = range(m.n_terms)
    m.K = range(m.k)
    m.var_values = pyo.Var(m.I, domain=pyo.Binary)
    
    
    """
    term_items holds the values of the literals in each term.
    """
    m.term_items = pyo.Var(m.J, m.K, domain=pyo.Binary)
    
    @m.Constraint( m.J, m.K )
    def evaluate_term_items(m: pyo.ConcreteModel, j: int, k: int) -> bool:
        if m.C[j, k] == 0:
            return m.term_items[j, k] == 0
        
        if m.C[j, k] > 0:
            return m.term_items[j, k] == m.var_values[m.C[j, k] - 1]
        
        if m.C[j, k] < 0:
            return m.term_items[j, k] == 1 - m.var_values[-m.C[j, k] - 1]
        
        assert False, "Unreachable"


    """
    term_values holds the values of the terms.
    
    Note: every term is a disjunction of literals
    """
    m.term_values = pyo.Var(m.J, domain=pyo.Binary)

    @m.Constraint(m.J, m.K)
    def evaluate_terms_low(m: pyo.ConcreteModel, j: int, k: int) -> bool:
        return m.term_values[j] >= m.term_items[j, k]
    
    @m.Constraint(m.J)
    def evaluate_terms_high(m: pyo.ConcreteModel, j: int) -> bool:
        return m.term_values[j] <= sum(m.term_items[j, ...])
    
    """
    The objective is to maximize the number of satisfied terms.
    """
    m.obj = pyo.Objective(expr=sum(m.term_values[...]), sense=pyo.maximize)
    
    return m

m: pyo.ConcreteModel = get_model(k, n_vars, n_terms, C)

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

pyo.value(m.obj)

m.var_values.extract_values()

m.term_values.extract_values()

[bool(pyo.value(m.var_values[i])) for i in range(m.n_vars)]

[[bool(pyo.value(m.term_items[i, j])) for j in range(m.k)] for i in range(m.n_terms)]

def gen_random_3sat(n_vars: int, n_terms: int) -> np.ndarray['n_terms,3']:
    """
    Returns C - the matrix of the CNF terms.
    A term is a disjunction of 3 literals.
    A literal is a var or -var.
    A var is an integer in [1, n_vars]
    """
    
    C: np.ndarray['n_terms,3'] = 1 + np.random.randint(n_vars, size=(n_terms, 3))
    C *= np.random.choice([-1, 1], size=(n_terms, 3))
    
    return C

gen_random_3sat(6, 10)

def convert_3sat_to_task(C: np.ndarray['n_terms,3']) -> tuple[int, int, int, np.ndarray['n_terms,3']]:
    n_vars: int = np.max(np.abs(C))
    n_terms: int = C.shape[0]
    k: int = 3
    assert C.shape == (n_terms, k)
    
    return k, n_vars, n_terms, C

def pysat_solve_3sat(C: np.ndarray['n_terms,3']) -> bool:
    with Solver() as solver:
        for i in range(C.shape[0]):
            clause = list(map(int, C[i, ...])) 
            # print(">", clause)
            solver.add_clause(clause)
        return solver.solve()

def my_solve_3sat(C: np.ndarray['n_terms,3']) -> bool:
    m: pyo.ConcreteModel = get_model(*convert_3sat_to_task(C))
    
    solver = pyo.SolverFactory('cbc')
    solver.solve(m)
    
    return np.isclose(pyo.value(m.obj), m.n_terms)

def test_3sat(n: int = 100) -> bool:
    for _ in range(n):
        C = gen_random_3sat(10, 20)
        
        matches: bool = pysat_solve_3sat(C) == my_solve_3sat(C)
        
        if not matches:
            print(f"Discrpeancy found for C = {C}!")
            return False
    
    return True

assert test_3sat()
print("Yay!")