#!python3

from pysat.formula import CNF
from pysat.solvers import Glucose4

import pathlib as pl
import random
import bz2
import sys
import argparse 

CORE_DIR = pl.Path("./cores")
EXTENDED_DIR = pl.Path("./extensions")


def main():
    args = parse_args()

    instances = get_instance_paths(pl.Path(args.instance_dir))
    
    CORE_DIR = pl.Path(args.cores_dir)
    EXTENDED_DIR = pl.Path(args.extensions_dir)
    
    for instance in instances:
        run(instance, args.nr_per_instance, args.verbosity > 0)

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument( "-i"
                        ,"--instance_dir"
                        ,type=str
                        ,default="./instances"
                        ,help="directory containing the instances")
    parser.add_argument( "-c"
                        ,"--cores_dir"
                        ,type=str
                        ,default="./cores"
                        ,help="directory containing the cores")
    parser.add_argument( "-e"
                        ,"--extensions_dir"
                        ,type=str
                        ,default="./extensions"
                        ,help="directory target directory to store the generated core extensions")
    parser.add_argument( "-n"
                        ,"--nr_per_instance"
                        ,type=int
                        ,default=1
                        ,help="number of extensions per instance")
    parser.add_argument( "-v"
                        ,"--verbosity"
                        ,type=int
                        ,default=1
                        ,help="1 to print progress, 0 for silent mode")

    return parser.parse_args()

def get_instance_paths(instances_dir):
    instances = []

    for p in instances_dir.iterdir():
        if str(p).endswith(".cnf.bz2"):
            instances.append(p)

    return instances 

def run(instance, calls_per_instance, verb):
    ipath = pl.Path(instance)
    iname = ipath.stem[:-4]

    f = CNF(str(ipath))
    
    cores = parse_cores(iname)
    

    for i in range(calls_per_instance):
        extension = gen_extension(cores)
        
        if verb and i % 20 == 0:
            print("instance: {}".format(iname))

        if verb:
            print("extension: {} | #cores: {} ({:2.2f}%)".format(i+1, len(extension), 100 * len(extension) / len(cores)))

        ext_f = CNF()
        ext_f.extend(extension)
        
        fpath, nr = next_file_path(iname)
        
        if not EXTENDED_DIR.exists():
            EXTENDED_DIR.mkdir(parents=True)

        ext_f.to_file(fpath,
                      comments=["c extending:{} nr:{}".format(iname, nr),
                                "c number of cores: {}".format(len(extension))],
                      compress_with="bzip2")

    if verb:
        print()

def gen_extension(cores):

    p = (len(cores) / 10 ) / len(cores)
   
    extension = []
    while len(extension) == 0:
        for core in cores:
            if p >= random.random():
                extension.append(core)


    return extension

def parse_cores(instance):
    cores = []
    
    with bz2.open(CORE_DIR / (instance + ".cores.bz2"), "tr") as cores_file:
        cores = list(map(lambda l: l.strip(), cores_file.readlines()))
    
    cores = list(map(lambda l: l.split(","), cores))
    
    for i in range(len(cores)):
        cores[i] = list(map(int, cores[i]))
        cores[i] = list(map(lambda lit: -1 * lit, cores[i]))

    return cores

def next_file_path(name):
    counter = 0

    ext_file_path = EXTENDED_DIR / ("{}_{:04}.cnf.bz2".format(name, counter))

    while ext_file_path.exists():
        counter += 1
        ext_file_path = EXTENDED_DIR / ("{}_{:04}.cnf.bz2".format(name, counter))
        
    return ext_file_path, counter 

if __name__ == "__main__":
    main()