
import os
import argparse

def add_weight_annotations(cnf_file, verbose=False):
    with open(cnf_file, 'r') as f:
        lines = f.readlines()

    # Find the number of variables from the CNF header
    for line in lines:
        if line.startswith('p cnf'):
            parts = line.split()
            nvars = int(parts[2])
            break

    # Create the output file with the same base name but with .sharpsat_td extension
    base_name = os.path.splitext(cnf_file)[0]
    output_file = base_name + '.sharpsat_td'

    with open(output_file, 'w') as f:
        # Write the original CNF content
        for line in lines:
            f.write(line)

        # Write the weight annotations
        for idx in range(1, nvars + 1):
            f.write(f"c p weight {idx} {idx} 0\n")
            f.write(f"c p weight {-idx} {-idx} 0\n")

    if verbose:
        print(f"Output written to {output_file}")

def main():
    parser = argparse.ArgumentParser(description="Add weight annotations to CNF files for SharpSAT-TD.")
    parser.add_argument("cnf_file", help="Path to the input CNF file")
    parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose output")

    args = parser.parse_args()

    cnf_file = args.cnf_file

    if not os.path.exists(cnf_file):
        raise FileNotFoundError(f"File {cnf_file} not found")

    add_weight_annotations(cnf_file, verbose=args.verbose),

if __name__ == "__main__":
    main()
