#!python3

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

import pathlib as pl
import random
import bz2
import sys

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

usage_str = "usage: gen_eq_instance.py [instance dir] [num per instance]"

def main():
    instances_dir = None

    try:
        instances_dir = pl.Path(sys.argv[1])
    except:
        print(usage_str)
        return
    else:
        if not instances_dir.is_dir():
            print(usage_str)
            return

    num_per_instance = 0
    try:
        num_per_instance = int(sys.argv[2])
    except:
        print(usage_str)
        return


    instances = get_instance_paths(instances_dir)
    
    for instance in instances:
        run(instance, num_per_instance)

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):
    ipath = pl.Path(instance)
    iname = ipath.stem[:-4]

    f = CNF(str(ipath))
    
    for i in range(calls_per_instance):
        extension = gen_extension(iname, len(f.clauses) / 10)
        
        ext_f = CNF()
        ext_f.extend(extension)

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

def gen_extension(instance, expected_num_of_core):
    cores = parse_cores(instance)

    p = expected_num_of_core / len(cores)
    
    extension = []
    
    print((expected_num_of_core * 10) / len(cores))
    print(expected_num_of_core * 10, len(cores))

    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()