#!/usr/bin/env python3

from util import kSAT
from util import SAT2QUBO
from util.SATquboResult import SATquboResult
import argparse
from dwave_qbsolv import QBSolv
import os
import collections
import json
from tqdm import tqdm
from util import script as scriptUtils

def main():
    arguments = __parseArguments()
    
    satInstance = kSAT.kSAT()
    
    print("reading ksat...")
    satInstance.readDIMACS(arguments["instancePath"])
    
    print()
    result = __runWMISquboOnSatInstance(satInstance)
    
    resultPath = os.path.join(
                 os.path.join(arguments["resultDir"],
                              "run%d" % arguments["run"]),
                              "%s.out" % arguments["instanceFileName"])
    
    print()
    print("writing results to file...")
    __writeResult(result, resultPath)

def __parseArguments():
    parser = scriptUtils.ArgParser()
    
    parser.addInstanceDirArg()
    parser.addArg(alias="instancePath", shortFlag="i", longFlag="instance",
                  help="instance file, has to be in DIMACS format", type=str)
    
    parser.addArg(alias="resultDir", shortFlag="o", longFlag="wmis_result_dir",
                  help="the wmis result directory", type=str,
                  ignoreDatabaseConfig=False)
    
    parser.addArg(alias="run", shortFlag="r", longFlag="run",
                  help="results will get saved unter [instance]_[run].out", type=int)
                                      
    arguments = parser.parse()
    
    arguments["instanceFileName"] = os.path.basename(arguments["instancePath"])
    
    return arguments
        
        
        
def __runWMISquboOnSatInstance(satInstance):
    print("generating wmis qubo...")
    qubo = SAT2QUBO.WMISdictQUBO(satInstance)
    
    print()
    print("running gbsolv...")
    qbresult = QBSolv().sample_qubo(Q=qubo, find_max=True)
    
    print()
    print("packing results...")
    results = __packResults(satInstance, qbresult)
    
    return results
    
def __packResults(satInstance, qbresult):
    results = []
    
    samples = list(qbresult.samples())
    occurrences = qbresult.data_vectors["num_occurrences"]
    
    for i in tqdm(range(len(samples))):
        quboResult = __satQuboResultFromSample(samples[i])
        quboResult.setOccurrences(occurrences[i])
        quboResult.setSatisfiesInstance(satInstance)
        
        results.append(quboResult)
    
    return results
    
    
def __writeResult(results, resultPath):
    resultDir = os.path.dirname(resultPath)
    
    if not os.path.exists(resultDir):
        os.makedirs(resultDir)
    
    resultFile = open(resultPath, "w+")
    
    for result in tqdm(results):
        
        resultFile.write(json.dumps(result.toPrimitive()))
        resultFile.write("\n\n")
        
    resultFile.close()

def __satQuboResultFromSample(sample):
    result = SATquboResult()
    
    for binding in sample:
        isActive = True if sample[binding] == 1 else False
        
        result.addBinding(binding, isActive)
        #if sample[binding] == 1:
            #result.addActiveBinding(binding)
    
    return result
    

if __name__ == "__main__":
    main()