import numpy as np
import matplotlib.pyplot as plt
from docplex.mp.model import Model

from qiskit import BasicAer
from qiskit.algorithms import QAOA, NumPyMinimumEigensolver
from qiskit_optimization.algorithms import CplexOptimizer, MinimumEigenOptimizer
from qiskit_optimization.algorithms.admm_optimizer import ADMMParameters, ADMMOptimizer
from qiskit_optimization import QuadraticProgram

from qiskit_optimization.converters import InequalityToEquality, IntegerToBinary, LinearEqualityToPenalty

def data_bins(results, wj, n, m, l=0, simplify=False):
    """save the results on a dictionary with the three items, bins, items, and index.
    results (cplex.solve): results of the optimization
    wj: (array (1,m): weights of the items
    n: (int) number of items
    m: (int) number of bins
    """
    if simplify:
        bins = np.ones((m,))
        if m-l > 0: 
            bins[m-l-1:m] = results[:m-l]
        items = np.zeros((m,n))
        items[:,1:] =  results[m-l:(m-1)*n+m-l].reshape(m,n-1)
        items[0,0] = 1
        items = items.reshape(m,n) * wj
        return {"bins":bins, "items":items,"index":np.arange(m)}
    else:        
        return {"bins":results[:m], "items":results[m:m+m*n].reshape(m,n) * wj, "index":np.arange(m)}
def plot_bins(results, wj, n, m, l=0,simplify=False):
    """plot in a bar diagram the results of an optimization bin packing problem"""
    res = data_bins(results.x, wj, n, m, l, simplify)
    plt.figure()
    ind = res["index"]
    plt.bar(ind, res["items"][:,0], label=f"item {0}")
    suma = bottom=res["items"][:,0]
    for j in range(1,n):
        plt.bar(ind, res["items"][:,j], bottom=suma, label=f"item {j}")
        suma += res["items"][:,j]
    plt.hlines(Q,0-0.5,m-0.5,linestyle="--", color="r",label="Max W")
    plt.xticks(ind)
    plt.xlabel("Bin")
    plt.ylabel("Weight")
    plt.legend()

np.random.seed(2)
n = 3 # number of bins
m = n # number of items
Q = 40 # max weight of a bin

wj = np.random.randint(1,Q,n) # Randomly picking the item weight

# Construct model using docplex
mdl = Model("BinPacking")

x = mdl.binary_var_list([f"x{i}" for i in range(n)]) # list of variables that represent the bins
e =  mdl.binary_var_list([f"e{i//m},{i%m}" for i in range(n*m)]) # variables that represent the items on the specific bin

objective = mdl.sum([x[i] for i in range(n)])

mdl.minimize(objective)

for j in range(m):
    # First set of constraints: the items must be in any bin
    constraint0 = mdl.sum([e[i*m+j] for i in range(n)])
    mdl.add_constraint(constraint0 == 1, f"cons0,{j}")
    
for i in range(n):
    # Second set of constraints: weight constraints
    constraint1 = mdl.sum([wj[j] * e[i*m+j] for j in range(m)])
    mdl.add_constraint(constraint1 <= Q * x[i], f"cons1,{i}")


# Load quadratic program from docplex model
qp = QuadraticProgram()
qp.from_docplex(mdl)
print(qp.export_as_lp_string())

# convert from DOcplex model to Qiskit Quadratic program
qp = QuadraticProgram()
qp.from_docplex(mdl)

# Solving Quadratic Program using CPLEX
cplex = CplexOptimizer()
result = cplex.solve(qp)
print(result)
plot_bins(result, wj, n, m)

# Construct model using docplex
mdl = Model("BinPacking_simplify")

l = int(np.ceil(np.sum(wj)/Q))
x = mdl.binary_var_list([f"x{i}" for i in range(m)]) # list of variables that represent the bins
e =  mdl.binary_var_list([f"e{i//m},{i%m}" for i in range(n*m)]) # variables that represent the items on the specific bin

objective = mdl.sum([x[i] for i in range(n)])

mdl.minimize(objective)

for j in range(m):
    # First set of constraints: the items must be in any bin
    constraint0 = mdl.sum([e[i*m+j] for i in range(n)])
    mdl.add_constraint(constraint0 == 1, f"cons0,{j}")
    
for i in range(n):
    # Second set of constraints: weight constraints
    constraint1 = mdl.sum([wj[j] * e[i*m+j] for j in range(m)])
    mdl.add_constraint(constraint1 <= Q * x[i], f"cons1,{i}")


# Load quadratic program from docplex model
qp = QuadraticProgram()
qp.from_docplex(mdl)
# Simplifying the problem
for i in range(l):
    qp = qp.substitute_variables({f"x{i}":1}) 
qp = qp.substitute_variables({"e0,0":1}) 
for i in range(1,m):
    qp = qp.substitute_variables({f"e{i},0":0})
print(qp.export_as_lp_string())

simplify_result = cplex.solve(qp)
print(simplify_result)
plot_bins(simplify_result, wj, n, m, l, simplify=True)

ineq2eq = InequalityToEquality()
qp_eq = ineq2eq.convert(qp)
print(qp_eq.export_as_lp_string())
print(f"The number of variables is {qp_eq.get_num_vars()}")

int2bin = IntegerToBinary()
qp_eq_bin = int2bin.convert(qp_eq)
print(qp_eq_bin.export_as_lp_string())
print(f"The number of variables is {qp_eq_bin.get_num_vars()}")

lineq2penalty = LinearEqualityToPenalty()
qubo = lineq2penalty.convert(qp_eq_bin)
print(f"The number of variables is {qp_eq_bin.get_num_vars()}")
print(qubo.export_as_lp_string())

result = cplex.solve(qubo)
print(result)

data_bins(result.x, wj, n, m, l=l, simplify=True)
plot_bins(result, wj, n, m, l=l, simplify=True)

from qiskit import Aer
backend = Aer.get_backend("qasm_simulator")

qaoa = MinimumEigenOptimizer(QAOA(reps=3, quantum_instance=backend))
result_qaoa = qaoa.solve(qubo)
print(result_qaoa)

plot_bins(result, wj, n, m, l, simplify=True)
plt.title("QAOA solution", fontsize=18)

# Proposal
plt.figure()
alpha = -2
x = np.linspace(-5,5,100)
f = (-(x+alpha) + (x+alpha)**2)
plt.plot(x,f)
plt.grid()
plt.xlabel("x")
plt.ylabel("f(x)")

# Construct model using docplex
mdl = Model("binPackingMyApproach")

x = mdl.binary_var_list([f"x{i}" for i in range(n)]) # list of variables that represent the bins
e =  mdl.binary_var_list([f"e{i//m},{i%n}" for i in range(n*m)]) # variables that represent the items on the specific bin

objective = mdl.sum([x[i] for i in range(m)])
cons = 0
alpha = -10
for i in range(m):
    cons_1 = 0
    cons_1 += Q*x[i]
    for j in range(n):
        cons_1 -= wj[j]*e[i*m+j]
    cons += (-(cons_1+alpha) + (cons_1+alpha)**2)

mdl.minimize(objective + cons)

for j in range(m):
    # First set of constraints: the items must be in any bin
    constraint0 = mdl.sum([e[i*m+j] for i in range(n)])
    mdl.add_constraint(constraint0 == 1, f"cons0,{j}")
    
# # Load quadratic program from docplex model
qp = QuadraticProgram()
qp.from_docplex(mdl)
# Simplifying the problem
for i in range(l):
    qp = qp.substitute_variables({f"x{i}":1}) 
qp = qp.substitute_variables({"e0,0":1}) 
for i in range(1,m):
    qp = qp.substitute_variables({f"e{i},0":0})
print(qp.export_as_lp_string())
print(f"The number of variables is {qp.get_num_vars()}")

# Solving Quadratic Program using CPLEX
cplex = CplexOptimizer()
result = cplex.solve(qp)
print(result)
plot_bins(result, wj, n, m, l, simplify=True)


qubo_myapp = lineq2penalty.convert(qp)
result_myapp = qaoa.solve(qubo_myapp)
print(result_myapp)
plot_bins(result_myapp, wj, n, m, l, simplify=True)
plt.title("QAOA solution")

import qiskit.tools.jupyter
%qiskit_version_table
%qiskit_copyright

