#!/bin/python3
import sys,time,re,os,copy, ast
from pprint import pprint
sys.path.insert(0, '../../')
from fred import *
from z3 import *

start_time = time.time()

if sys.argv[1] == 'test':
    input_f = 'test'
elif sys.argv[1] == 'input':
    input_f = 'input'
else:
    print('No argv provided')
    exit()

def parse_line(line):
    pattern = re.compile(r'^\[(.*)\]\s(.*)\s{1}\{(.*)\}$')

    match = re.match(pattern,line)
    lights = list2int(list(match.group(1).replace('.','0').replace('#','1')))
    buttons = [ [v] if isinstance((v := ast.literal_eval(b)), int) else list(v) for b in match.group(2).split(' ') ]
    joltage = list2int(list(match.group(3).split(',')))
    return lights,joltage,buttons

#########################################
#                                       #
#              Part 1                   #
#                                       #
#########################################

def solve_part1_z3(target, buttons):
    """Solve using Z3 with mod 2 arithmetic."""
    n_lights = len(target)
    n_buttons = len(buttons)
    
    # Variables: how many times to press each button (0 or 1 for mod 2)
    presses = [Int(f'p{i}') for i in range(n_buttons)]
    
    opt = Optimize()
    
    # Each press is 0 or 1 (mod 2, pressing twice = 0 effect)
    for p in presses:
        opt.add(Or(p == 0, p == 1))
    
    # For each light: sum of button effects mod 2 = target
    for light_idx in range(n_lights):
        total = 0
        for btn_idx, btn in enumerate(buttons):
            if light_idx in btn:
                total += presses[btn_idx]
        opt.add(total % 2 == target[light_idx])
    
    # Minimize total presses
    opt.minimize(Sum(presses))
    
    if opt.check() == sat:
        m = opt.model()
        return sum(m[p].as_long() for p in presses)
    return None


def part1():
    score = 0
    lines = loadFile(input_f)
    
    for line in lines:
        target, _, buttons = parse_line(line)
        presses = solve_part1_z3(target, buttons)
        if presses is None:
            return None
        score += presses
    
    return score



start_time = time.time()
p1 = part1()
print('Part 1:', p1, '\t\t', round((time.time() - start_time)*1000), 'ms')

#########################################
#                                       #
#              Part 2                   #
#                                       #
#########################################

def solve_part2_z3(target, buttons):
    """Solve using Z3 with integer arithmetic."""
    n_counters = len(target)
    n_buttons = len(buttons)
    
    # Variables: how many times to press each button (non-negative integer)
    presses = [Int(f'p{i}') for i in range(n_buttons)]
    
    opt = Optimize()
    
    # Each press must be non-negative
    for p in presses:
        opt.add(p >= 0)
    
    # For each counter: sum of button effects = target
    for counter_idx in range(n_counters):
        total = 0
        for btn_idx, btn in enumerate(buttons):
            if counter_idx in btn:
                total += presses[btn_idx]
        opt.add(total == target[counter_idx])
    
    # Minimize total presses
    opt.minimize(Sum(presses))
    
    if opt.check() == sat:
        m = opt.model()
        return sum(m[p].as_long() for p in presses)
    return None


def part2():
    score = 0
    lines = loadFile(input_f)
    
    for line_num, line in enumerate(lines, 1):
        _, target, buttons = parse_line(line)
        presses = solve_part2_z3(target, buttons)
        if presses is None:
            print(f"Part 2 Machine {line_num} cannot be configured!")
            return None
        score += presses
    
    return score

start_time = time.time()
p2 = part2()
print('Part 2:', p2, '\t\t', round((time.time() - start_time)*1000), 'ms')