#!/usr/bin/env python3

import numpy as np
from . import kSAT
from tqdm import tqdm
import math

__VERTEX_WEIGHT__ = -1
__EDGE_WEIGHT__ = 2

def WMISdictQUBO(kSATInstance):
    quboInstance = {}
    
    for clauseIndex in range(kSATInstance.getNumberOfClauses()):
        clause = kSATInstance.getClause(clauseIndex)
        
        # build triangles
        for varIndexInClause in range(len(clause)):
            label = kSATInstance.getLableOfBinding(clauseIndex,
                                                   clause[varIndexInClause])
            quboInstance[(label, label)] = __VERTEX_WEIGHT__
            
            for i in range(varIndexInClause + 1, len(clause)):
                targetLabel = kSATInstance.getLableOfBinding(clauseIndex,
                                                             clause[i])
                
                quboInstance[(label, targetLabel)] = __EDGE_WEIGHT__
        
    # connect conflicts
    for conflict in kSATInstance.getConflicts():
        quboInstance[(conflict[0], conflict[1])] = __EDGE_WEIGHT__
    
    return quboInstance

# only 3sat
def primitiveQUBO(sat):
    quboInstance = {}
    
    chains = {}
    
    
    for clauseIndex in range(sat.getNumberOfClauses()):
        clause = sat.getClause(clauseIndex);
        
        #build clause primitives
        
        lit1 = "c{}_l{}".format(clauseIndex, clause[0])
        lit2 = "c{}_l{}".format(clauseIndex, clause[1])
        lit3 = "c{}_l{}".format(clauseIndex, clause[2])
        aux1 = "a{}_{}".format(clauseIndex, 1)
        aux2 = "a{}_{}".format(clauseIndex, 2)
        aux3 = "a{}_{}".format(clauseIndex, 3)
        aux4 = "a{}_{}".format(clauseIndex, 4)
        
        quboInstance[(lit1, lit1)] = 1;
        quboInstance[(lit2, lit2)] = 1;
        quboInstance[(lit3, lit3)] = 1;
        quboInstance[(aux1, aux1)] = -2;
        quboInstance[(aux2, aux2)] = 1;
        quboInstance[(aux3, aux3)] = -2;
        quboInstance[(aux4, aux4)] = -2;
        
        quboInstance[(lit1, lit2)] = 1;
        quboInstance[(lit1, aux1)] = -2;
        quboInstance[(lit2, aux1)] = -2;
        quboInstance[(aux1, aux2)] = -2;
        quboInstance[(aux2, aux3)] = -2;
        quboInstance[(aux2, lit3)] = 1;
        quboInstance[(lit3, aux3)] = -2;
        quboInstance[(aux3, aux4)] = -2;
        
        if clause[0] in chains:
            chains[clause[0]].append(lit1)
        else:
            chains[clause[0]] = [lit1]
            
        if clause[1] in chains:
            chains[clause[1]].append(lit2)
        else:
            chains[clause[1]] = [lit2]
            
        if clause[2] in chains:
            chains[clause[2]].append(lit3)
        else:
            chains[clause[2]] = [lit3]
    
    #build chains
    longestChain = 0;
    for lit, nodes in chains.items():
        if len(nodes) > longestChain:
            longestChain = len(nodes)
            
        if lit > 0 and -1 * lit in chains:
            quboInstance[(chains[lit][0], chains[-1*lit][0])] = 2
            
    print("longest chain = {}".format(longestChain))
    
    for nodes in chains.values():
        while len(nodes) > 1:
            quboInstance[(nodes[0], nodes[1])] = -2;
            nodes.pop(0)
    
    return quboInstance
    
class QuboWriter:
    def __init__(self, qubo):
        self.__labelIndexDict = {}
        self.__qubo = qubo
        
        self.__initLabelIndexDict(self.__qubo)
        
        print(self.__labelIndexDict)
    
    def __initLabelIndexDict(self, qubo):
        indexCount = 0
        
        for coupler in qubo:
            
            label1 = coupler[0]
            label2 = coupler[1]
            
            if label1 not in self.__labelIndexDict:
                self.__labelIndexDict[label1] = indexCount
                indexCount += 1
            
            if label2 not in self.__labelIndexDict:
                self.__labelIndexDict[label2] = indexCount
                indexCount += 1
    
    def write(self, filePath):
        quboFile = open(filePath, "w")
    
        self.__writeQuboFileHeader(quboFile)
        self.__writeQuboFileNodes(quboFile)
        self.__writeQuboFileCouplers(quboFile)
    
        quboFile.close()
    
    def __writeQuboFileHeader(self, quboFile):
        numberOfNodes = len(self.__labelIndexDict)
        numberOfCouplers = len(self.__qubo) - numberOfNodes
    
        quboFile.write("c\n")
        quboFile.write("c this is a generated qubo file\n")
        quboFile.write("c\n")
    
        quboFile.write("p qubo 0 %d %d %d\n" %(numberOfNodes,
                                               numberOfNodes,
                                               numberOfCouplers))
                                               
    def __writeQuboFileNodes(self, quboFile):
        for label in self.__labelIndexDict:
            self.__writeCoupler(quboFile, (label, label))
            
    def __writeCoupler(self, quboFile, coupler):
            indices = self.__getNodeIndicesFromCoupler(coupler)
    
            quboFile.write("%d %d  %d\n" % (indices[0],
                                            indices[1],
                                            self.__qubo[coupler]))

    def __getNodeIndicesFromCoupler(self, coupler):
        index1 = self.__labelIndexDict[coupler[0]]
        index2 = self.__labelIndexDict[coupler[1]]
        
        if index1 <= index2:
            return [index1, index2]
        else:
            return [index2, index1]
            
    def __writeQuboFileCouplers(self, quboFile):
        for coupler in self.__qubo:
            if coupler[0] != coupler[1]:
                self.__writeCoupler(quboFile, coupler)