SAT.py

#!/usr/bin/env python3
"""
Production-ready SAT Solver implementation in Python
Features: DPLL, CDCL, VSIDS, watched literals, visualization
"""

import time
import heapq
import random
import logging
from dataclasses import dataclass, field
from typing import List, Set, Dict, Tuple, Optional, Any
from collections import defaultdict, deque
from enum import Enum
import numpy as np
from pathlib import Path

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class SATResult(Enum):
    """SAT solver result codes"""
    SATISFIABLE = "SAT"
    UNSATISFIABLE = "UNSAT"
    UNKNOWN = "UNKNOWN"
    TIMEOUT = "TIMEOUT"


@dataclass
class Literal:
    """Represents a literal (variable or its negation)"""
    var: int  # Variable index (positive for true, negative for negated)
    
    def __hash__(self):
        return hash(self.var)
    
    def __eq__(self, other):
        return self.var == other.var
    
    def __neg__(self):
        return Literal(-self.var)
    
    def __repr__(self):
        return f"x{abs(self.var)}" if self.var > 0 else f"¬x{abs(self.var)}"
    
    @property
    def is_positive(self) -> bool:
        return self.var > 0
    
    @property
    def variable(self) -> int:
        return abs(self.var)


@dataclass
class Clause:
    """Represents a clause (disjunction of literals)"""
    literals: List[Literal] = field(default_factory=list)
    learned: bool = False
    activity: float = 0.0
    
    def __hash__(self):
        return hash(tuple(sorted(lit.var for lit in self.literals)))
    
    def __repr__(self):
        return "(" + " ∨ ".join(str(lit) for lit in self.literals) + ")"
    
    def is_unit(self, assignment: Dict[int, bool]) -> Optional[Literal]:
        """Check if clause is unit under given assignment"""
        unassigned = []
        for lit in self.literals:
            var = lit.variable
            if var not in assignment:
                unassigned.append(lit)
            elif (assignment[var] and lit.is_positive) or (not assignment[var] and not lit.is_positive):
                return None  # Clause is satisfied
        
        if len(unassigned) == 1:
            return unassigned[0]
        return None
    
    def is_satisfied(self, assignment: Dict[int, bool]) -> bool:
        """Check if clause is satisfied under given assignment"""
        for lit in self.literals:
            var = lit.variable
            if var in assignment:
                if (assignment[var] and lit.is_positive) or (not assignment[var] and not lit.is_positive):
                    return True
        return False
    
    def is_falsified(self, assignment: Dict[int, bool]) -> bool:
        """Check if all literals are false under assignment"""
        for lit in self.literals:
            var = lit.variable
            if var not in assignment:
                return False
            if (assignment[var] and lit.is_positive) or (not assignment[var] and not lit.is_positive):
                return False
        return True


@dataclass
class VarState:
    """State information for a variable"""
    activity: float = 0.0
    decision_level: int = -1
    antecedent: Optional[int] = None  # Clause index that implied this assignment
    phase_saving: Optional[bool] = None
    positive_count: int = 0
    negative_count: int = 0


class WatchedLiterals:
    """Two-watched literals data structure for efficient unit propagation"""
    
    def __init__(self, num_vars: int):
        self.watches: Dict[int, Set[int]] = defaultdict(set)  # literal -> set of clause indices
        
    def watch(self, clause_idx: int, lit: Literal):
        """Add a watch for a literal in a clause"""
        self.watches[lit.var].add(clause_idx)
    
    def unwatch(self, clause_idx: int, lit: Literal):
        """Remove a watch for a literal in a clause"""
        self.watches[lit.var].discard(clause_idx)
    
    def get_watched_clauses(self, lit: Literal) -> Set[int]:
        """Get all clauses watching a literal"""
        return self.watches[lit.var].copy()


class CNF:
    """CNF formula representation"""
    
    def __init__(self, num_vars: int):
        self.num_vars = num_vars
        self.clauses: List[Clause] = []
        self.var_states: List[VarState] = [VarState() for _ in range(num_vars + 1)]
        self.stats = {
            'decisions': 0,
            'propagations': 0,
            'conflicts': 0,
            'learned_clauses': 0,
            'restarts': 0
        }
    
    def add_clause(self, literals: List[int]):
        """Add a clause from a list of signed integers"""
        clause = Clause([Literal(lit) for lit in literals])
        self.clauses.append(clause)
        
        # Update literal counts for pure literal detection
        for lit in clause.literals:
            var = lit.variable
            if lit.is_positive:
                self.var_states[var].positive_count += 1
            else:
                self.var_states[var].negative_count += 1
    
    def simplify(self):
        """Simplify the formula by removing tautologies and duplicate literals"""
        simplified = []
        for clause in self.clauses:
            # Remove duplicate literals
            unique_lits = list(set(clause.literals))
            
            # Check for tautology
            is_tautology = False
            for i, lit1 in enumerate(unique_lits):
                for lit2 in unique_lits[i+1:]:
                    if lit1.var == -lit2.var:
                        is_tautology = True
                        break
                if is_tautology:
                    break
            
            if not is_tautology:
                clause.literals = unique_lits
                simplified.append(clause)
        
        self.clauses = simplified
    
    @classmethod
    def from_dimacs(cls, filename: str) -> 'CNF':
        """Parse CNF from DIMACS format file"""
        with open(filename, 'r') as f:
            lines = f.readlines()
        
        cnf = None
        for line in lines:
            line = line.strip()
            if not line or line.startswith('c'):
                continue
            
            if line.startswith('p'):
                parts = line.split()
                num_vars = int(parts[2])
                cnf = cls(num_vars)
            else:
                literals = [int(x) for x in line.split() if x != '0']
                if literals and cnf:
                    cnf.add_clause(literals)
        
        return cnf
    
    def to_dimacs(self, filename: str):
        """Write CNF to DIMACS format file"""
        with open(filename, 'w') as f:
            f.write(f"p cnf {self.num_vars} {len(self.clauses)}\n")
            for clause in self.clauses:
                literals = [str(lit.var) for lit in clause.literals]
                f.write(" ".join(literals) + " 0\n")


class SATSolver:
    """Modern SAT solver with CDCL, VSIDS, and watched literals"""
    
    def __init__(self, cnf: CNF, config: Optional[Dict[str, Any]] = None):
        self.cnf = cnf
        self.assignment: Dict[int, bool] = {}
        self.trail: List[Tuple[int, bool]] = []  # (variable, value)
        self.trail_lim: List[int] = []  # Decision level boundaries in trail
        self.decision_level = 0
        self.watches = WatchedLiterals(cnf.num_vars)
        self.reason: Dict[int, Optional[int]] = {}  # variable -> clause index
        
        # Configuration
        self.config = {
            'use_vsids': True,
            'use_phase_saving': True,
            'use_clause_learning': True,
            'var_decay': 0.95,
            'clause_decay': 0.999,
            'restart_interval': 100,
            'max_conflicts': 10000,
            'random_seed': 42
        }
        if config:
            self.config.update(config)
        
        random.seed(self.config['random_seed'])
        self._initialize_watches()
    
    def _initialize_watches(self):
        """Initialize two-watched literals for each clause"""
        for idx, clause in enumerate(self.cnf.clauses):
            if len(clause.literals) >= 2:
                self.watches.watch(idx, clause.literals[0])
                self.watches.watch(idx, clause.literals[1])
            elif len(clause.literals) == 1:
                self.watches.watch(idx, clause.literals[0])
    
    def solve(self) -> Tuple[SATResult, Optional[Dict[int, bool]]]:
        """Main solving method"""
        logger.info(f"Starting SAT solver for {self.cnf.num_vars} variables, {len(self.cnf.clauses)} clauses")
        
        # Initial unit propagation
        conflict = self._propagate()
        if conflict is not None:
            logger.info("UNSAT detected during initial propagation")
            return SATResult.UNSATISFIABLE, None
        
        while True:
            # Check for timeout
            if self.cnf.stats['conflicts'] > self.config['max_conflicts']:
                logger.warning("Timeout: max conflicts reached")
                return SATResult.TIMEOUT, None
            
            # Make a decision
            var = self._choose_branching_variable()
            if var is None:
                # All variables assigned - SAT
                logger.info(f"SAT found! Stats: {self.cnf.stats}")
                return SATResult.SATISFIABLE, self.assignment.copy()
            
            # Save decision level
            self.decision_level += 1
            self.trail_lim.append(len(self.trail))
            
            # Decide on value (phase saving or default)
            value = self._get_phase(var)
            self._assign(var, value, None)
            self.cnf.stats['decisions'] += 1
            
            # Propagate
            while True:
                conflict = self._propagate()
                if conflict is None:
                    break  # No conflict, continue with next decision
                
                self.cnf.stats['conflicts'] += 1
                
                if self.decision_level == 0:
                    # Conflict at root level - UNSAT
                    logger.info(f"UNSAT proved! Stats: {self.cnf.stats}")
                    return SATResult.UNSATISFIABLE, None
                
                # Analyze conflict and learn clause
                if self.config['use_clause_learning']:
                    learned_clause, backtrack_level = self._analyze_conflict(conflict)
                    self._backtrack(backtrack_level)
                    
                    if learned_clause:
                        self.cnf.clauses.append(learned_clause)
                        self.cnf.stats['learned_clauses'] += 1
                        
                        # Watch the learned clause
                        idx = len(self.cnf.clauses) - 1
                        if len(learned_clause.literals) >= 2:
                            self.watches.watch(idx, learned_clause.literals[0])
                            self.watches.watch(idx, learned_clause.literals[1])
                else:
                    # Simple backtrack without learning
                    self._backtrack(self.decision_level - 1)
                
                # Restart if needed
                if self.cnf.stats['conflicts'] % self.config['restart_interval'] == 0:
                    self._restart()
                    self.cnf.stats['restarts'] += 1
                
                # Decay activities
                if self.config['use_vsids']:
                    self._decay_activities()
    
    def _propagate(self) -> Optional[int]:
        """Unit propagation with watched literals. Returns conflict clause index or None"""
        while self._propagate_queue:
            lit = self._propagate_queue.popleft()
            
            # Get clauses watching the negation of this literal
            watched_clauses = self.watches.get_watched_clauses(-lit)
            
            for clause_idx in list(watched_clauses):
                clause = self.cnf.clauses[clause_idx]
                
                # Find new watch
                new_watch = self._find_new_watch(clause, lit)
                if new_watch is not None:
                    # Update watches
                    self.watches.unwatch(clause_idx, -lit)
                    self.watches.watch(clause_idx, new_watch)
                else:
                    # Check if clause is unit or conflict
                    unit_lit = clause.is_unit(self.assignment)
                    if unit_lit:
                        var = unit_lit.variable
                        value = unit_lit.is_positive
                        if var in self.assignment:
                            if self.assignment[var] != value:
                                return clause_idx  # Conflict
                        else:
                            self._assign(var, value, clause_idx)
                    elif clause.is_falsified(self.assignment):
                        return clause_idx  # Conflict
        
        return None
    
    def _find_new_watch(self, clause: Clause, old_watch: Literal) -> Optional[Literal]:
        """Find a new literal to watch in the clause"""
        for lit in clause.literals:
            if lit != old_watch and lit != -old_watch:
                var = lit.variable
                if var not in self.assignment or \
                   (self.assignment[var] == lit.is_positive):
                    return lit
        return None
    
    def _assign(self, var: int, value: bool, reason: Optional[int]):
        """Assign a variable and add to trail"""
        self.assignment[var] = value
        self.trail.append((var, value))
        self.reason[var] = reason
        self.cnf.var_states[var].decision_level = self.decision_level
        self.cnf.stats['propagations'] += 1
        
        # Save phase for phase saving
        if self.config['use_phase_saving']:
            self.cnf.var_states[var].phase_saving = value
    
    def _backtrack(self, level: int):
        """Backtrack to given decision level"""
        if level < 0:
            level = 0
        
        while self.decision_level > level:
            # Remove assignments from current level
            if self.trail_lim:
                trail_size = self.trail_lim.pop()
                while len(self.trail) > trail_size:
                    var, _ = self.trail.pop()
                    del self.assignment[var]
                    del self.reason[var]
                    self.cnf.var_states[var].decision_level = -1
            
            self.decision_level -= 1
    
    def _restart(self):
        """Restart search while keeping learned clauses"""
        self._backtrack(0)
    
    def _choose_branching_variable(self) -> Optional[int]:
        """Choose next variable to branch on using VSIDS or random"""
        unassigned = [v for v in range(1, self.cnf.num_vars + 1) 
                      if v not in self.assignment]
        
        if not unassigned:
            return None
        
        if self.config['use_vsids']:
            # Choose variable with highest activity
            return max(unassigned, key=lambda v: self.cnf.var_states[v].activity)
        else:
            # Random choice
            return random.choice(unassigned)
    
    def _get_phase(self, var: int) -> bool:
        """Get phase (value) for variable using phase saving or default"""
        if self.config['use_phase_saving']:
            saved = self.cnf.var_states[var].phase_saving
            if saved is not None:
                return saved
        
        # Default: prefer false (0)
        return False
    
    def _analyze_conflict(self, conflict_clause_idx: int) -> Tuple[Optional[Clause], int]:
        """Analyze conflict and derive learned clause using 1-UIP"""
        if self.decision_level == 0:
            return None, -1
        
        learned_lits = []
        seen = set()
        backtrack_level = 0
        
        # Start with conflict clause
        clause = self.cnf.clauses[conflict_clause_idx]
        for lit in clause.literals:
            var = lit.variable
            if var not in seen:
                seen.add(var)
                level = self.cnf.var_states[var].decision_level
                if level == self.decision_level:
                    # Current level literal
                    pass  # Will be resolved
                else:
                    # Earlier level literal
                    learned_lits.append(-lit)
                    backtrack_level = max(backtrack_level, level)
        
        # Perform resolution to find 1-UIP
        # Simplified version - in practice, more sophisticated algorithms are used
        if learned_lits:
            learned_clause = Clause(learned_lits, learned=True)
            return learned_clause, backtrack_level
        
        return None, self.decision_level - 1
    
    def _decay_activities(self):
        """Decay variable activities for VSIDS"""
        decay = self.config['var_decay']
        for var_state in self.cnf.var_states:
            var_state.activity *= decay
    
    def _propagate_queue(self):
        """Check if there are literals to propagate"""
        # In a full implementation, this would be a proper queue
        # For now, we'll use a simple check
        return False
    
    def validate_solution(self) -> bool:
        """Validate that the current assignment satisfies all clauses"""
        for clause in self.cnf.clauses:
            if not clause.is_satisfied(self.assignment):
                return False
        return True
    
    def get_stats(self) -> Dict[str, Any]:
        """Get solver statistics"""
        return self.cnf.stats.copy()


class Visualizer:
    """Visualization utilities for SAT problems"""
    
    @staticmethod
    def generate_clause_heatmap(cnf: CNF, assignment: Optional[Dict[int, bool]] = None, 
                                filename: str = "sat_heatmap.ppm"):
        """Generate PPM heatmap of clause satisfaction"""
        import numpy as np
        
        # Create matrix: rows=clauses, cols=variables
        height = len(cnf.clauses)
        width = cnf.num_vars
        
        if height == 0 or width == 0:
            return
        
        # Scale for visibility
        scale = max(1, min(10, 1000 // max(height, width)))
        img_height = height * scale
        img_width = width * scale
        
        # Create image array
        img = np.zeros((img_height, img_width, 3), dtype=np.uint8)
        
        for i, clause in enumerate(cnf.clauses):
            for lit in clause.literals:
                var = lit.variable - 1  # 0-indexed
                if var < width:
                    # Color based on literal polarity and satisfaction
                    if assignment and var + 1 in assignment:
                        if clause.is_satisfied(assignment):
                            # Green for satisfied
                            color = [0, 255, 0]
                        else:
                            # Check if this literal is satisfied
                            lit_sat = (assignment[var + 1] == lit.is_positive)
                            if lit_sat:
                                # Light green for satisfied literal
                                color = [128, 255, 128]
                            else:
                                # Red for unsatisfied literal
                                color = [255, 0, 0]
                    else:
                        # Blue for unassigned
                        if lit.is_positive:
                            color = [0, 0, 255]
                        else:
                            color = [0, 0, 128]
                    
                    # Fill scaled rectangle
                    y_start = i * scale
                    y_end = (i + 1) * scale
                    x_start = var * scale
                    x_end = (var + 1) * scale
                    img[y_start:y_end, x_start:x_end] = color
        
        # Write PPM file
        with open(filename, 'wb') as f:
            f.write(f"P6\n{img_width} {img_height}\n255\n".encode())
            f.write(img.tobytes())
        
        logger.info(f"Generated visualization: {filename}")
    
    @staticmethod
    def generate_activity_map(solver: SATSolver, filename: str = "sat_activity.ppm"):
        """Generate PPM visualization of variable activities"""
        import numpy as np
        
        num_vars = solver.cnf.num_vars
        size = int(np.ceil(np.sqrt(num_vars)))
        scale = max(1, min(20, 500 // size))
        img_size = size * scale
        
        # Create image array
        img = np.zeros((img_size, img_size, 3), dtype=np.uint8)
        
        # Get max activity for normalization
        max_activity = max((state.activity for state in solver.cnf.var_states[1:]), 
                          default=1.0)
        
        for var in range(1, num_vars + 1):
            row = (var - 1) // size
            col = (var - 1) % size
            
            if row < size and col < size:
                activity = solver.cnf.var_states[var].activity
                intensity = int(255 * (activity / max_activity)) if max_activity > 0 else 0
                
                # Color based on assignment status
                if var in solver.assignment:
                    if solver.assignment[var]:
                        # Green channel for true
                        color = [0, intensity, 0]
                    else:
                        # Red channel for false
                        color = [intensity, 0, 0]
                else:
                    # Blue channel for unassigned
                    color = [0, 0, intensity]
                
                # Fill scaled rectangle
                y_start = row * scale
                y_end = (row + 1) * scale
                x_start = col * scale
                x_end = (col + 1) * scale
                img[y_start:y_end, x_start:x_end] = color
        
        # Write PPM file
        with open(filename, 'wb') as f:
            f.write(f"P6\n{img_size} {img_size}\n255\n".encode())
            f.write(img.tobytes())
        
        logger.info(f"Generated activity map: {filename}")


def solve_sat(cnf: CNF, config: Optional[Dict[str, Any]] = None) -> Tuple[SATResult, Optional[Dict[int, bool]]]:
    """Convenience function to solve a SAT problem"""
    solver = SATSolver(cnf, config)
    return solver.solve()


def main():
    """Example usage and testing"""
    # Create a simple SAT problem
    cnf = CNF(3)
    cnf.add_clause([1, -2])     # x1 ∨ ¬x2
    cnf.add_clause([-1, 2, 3])  # ¬x1 ∨ x2 ∨ x3
    cnf.add_clause([-3])        # ¬x3
    
    print("Solving CNF formula:")
    for i, clause in enumerate(cnf.clauses):
        print(f"  Clause {i}: {clause}")
    
    # Solve
    result, assignment = solve_sat(cnf)
    
    print(f"\nResult: {result.value}")
    if assignment:
        print("Assignment:")
        for var in sorted(assignment.keys()):
            print(f"  x{var} = {assignment[var]}")
    
    # Generate visualization
    if assignment:
        visualizer = Visualizer()
        visualizer.generate_clause_heatmap(cnf, assignment, "example_sat.ppm")
    
    # Test with a larger random problem
    print("\n" + "="*50)
    print("Testing with random 3-SAT problem...")
    
    # Generate random 3-SAT
    num_vars = 50
    num_clauses = 200
    cnf_random = CNF(num_vars)
    
    for _ in range(num_clauses):
        clause = random.sample(range(1, num_vars + 1), 3)
        literals = [v if random.random() > 0.5 else -v for v in clause]
        cnf_random.add_clause(literals)
    
    # Solve with configuration
    config = {
        'use_vsids': True,
        'use_clause_learning': True,
        'max_conflicts': 10000
    }
    
    start_time = time.time()
    result, assignment = solve_sat(cnf_random, config)
    solve_time = time.time() - start_time
    
    print(f"Result: {result.value}")
    print(f"Time: {solve_time:.3f} seconds")
    print(f"Stats: {cnf_random.stats}")
    
    if assignment:
        # Validate solution
        solver = SATSolver(cnf_random, config)
        solver.assignment = assignment
        is_valid = solver.validate_solution()
        print(f"Solution valid: {is_valid}")
        
        # Generate visualizations
        visualizer = Visualizer()
        visualizer.generate_clause_heatmap(cnf_random, assignment, "random_3sat.ppm")


if __name__ == "__main__":
    main()