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 QuadraticProgramToQubo
from qiskit.algorithms.optimizers import COBYLA

%matplotlib inline 

np.random.seed(1)
num_items = 3 # number of items
num_bins = num_items # maximum number of bins
max_weight = 15 # max weight of a bin

weights = np.random.randint(1, max_weight, num_items) # Randomly picking the item weight

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

y = mdl.binary_var_list(num_bins, name="y") # list of variables that represent the bins
x =  mdl.binary_var_matrix(num_items, num_bins, "x") # variables that represent the items on the specific bin

objective = mdl.sum(y)

mdl.minimize(objective)

for i in range(num_items):
    # First set of constraints: the items must be in any bin
    mdl.add_constraint(mdl.sum(x[i, j] for j in range(num_bins)) == 1)
    
for j in range(num_bins):
    # Second set of constraints: weight constraints
    mdl.add_constraint(mdl.sum(weights[i] * x[i, j] for i in range(num_items)) <= max_weight * y[j])

# Load quadratic program from docplex model
qp = QuadraticProgram()
qp.from_docplex(mdl)
qubo = QuadraticProgramToQubo().convert(qp)# Create a converter from quadratic program to qubo representation

print(f"The number of variables prior converting to QUBO is {qp.get_num_vars()}")
print(f"The number of variables after converting to QUBO is {qubo.get_num_vars()}")

plt.figure()
plt.bar(0, qp.get_num_vars(), label="Prior QUBO")
plt.bar(1, qubo.get_num_vars(), label="QUBO")
plt.xticks(range(2), 2 * [""])
plt.xlim((-1,2))
plt.legend()
plt.ylabel("Num. vars")
plt.title(f"Num. vars for {num_items} items")

def BinPacking(num_items, num_bins, weights, max_weight, simplification=False):
    # Construct model using docplex
    mdl = Model("BinPacking")

    y = mdl.binary_var_list(num_bins, name="y") # list of variables that represent the bins
    x =  mdl.binary_var_matrix(num_items, num_bins, "x") # variables that represent the items on the specific bin

    objective = mdl.sum(y)

    mdl.minimize(objective)

    for i in range(num_items):
        # First set of constraints: the items must be in any bin
        mdl.add_constraint(mdl.sum(x[i, j] for j in range(num_bins)) == 1)

    for j in range(num_bins):
        # Second set of constraints: weight constraints
        mdl.add_constraint(mdl.sum(weights[i] * x[i, j] for i in range(num_items)) <= max_weight * y[j])

    # Load quadratic program from docplex model
    qp = QuadraticProgram()
    qp.from_docplex(mdl)
    if simplification:
        l = int(np.ceil(np.sum(weights)/max_weight))
        qp = qp.substitute_variables({f"y_{_}":1 for _ in range(l)}) # First simplification 
        qp = qp.substitute_variables({"x_0_0":1}) # Assign the first item into the first bin
        qp = qp.substitute_variables({f"x_0_{_}":0 for _ in range(1, num_bins)}) # as the first item is in the first 
                                                                                #bin it couldn't be in the other bins
    qubo = QuadraticProgramToQubo().convert(qp)# Create a converter from quadratic program to qubo representation
    return qubo

qubo_simp = BinPacking(num_items, num_bins, weights, max_weight, simplification=True)
print(f"After the simplification, the number of variables is {qubo_simp.get_num_vars()}")
plt.figure()
plt.bar(0, qp.get_num_vars(), label="Prior QUBO")
plt.bar(1, qubo.get_num_vars(), label="QUBO")
plt.bar(2, qubo_simp.get_num_vars(), label="QUBO Simplifications")
plt.xticks(range(3), 3 * [""])
plt.xlim((-1,3))
plt.legend()
plt.ylabel("Num. vars")
plt.title(f"Num. vars for {num_items} items")

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

def interpret(results, weights, num_items, num_bins, simplify=False):
    """
    Save the results as a list of list where each sublist represent a bin
    and the sublist elements represent the items weights
    
    Args:
    results: results of the optimization
    weights (list): weights of the items
    num_items: (int) number of items
    num_bins: (int) number of bins
    """
    if simplify:
        l = int(np.ceil(np.sum(weights)/max_weight))
        bins = l * [1] + list(results[:num_bins - l])
        items = results[num_bins - l: (num_bins - l) + num_bins * (num_items - 1)].reshape(num_items - 1, num_bins)
        items_in_bins = [[i+1 for i in range(num_items-1) if bins[j] and items[i, j]] for j in range(num_bins)]
        items_in_bins[0].append(0)
    else:
        bins = results[:num_bins]
        items = results[num_bins:(num_bins + 1) * num_items].reshape((num_items, num_bins))
        items_in_bins = [[i for i in range(num_items) if bins[j] and items[i, j]] for j in range(num_bins)]
    return items_in_bins

def get_figure(items_in_bins, weights, max_weight, title=None):
    """Get plot of the solution of the Bin Packing Problem.

    Args:
        result : The calculated result of the problem

    Returns:
        fig: A plot of the solution, where x and y represent the bins and
        sum of the weights respectively.
    """
    colors = plt.cm.get_cmap("jet", len(weights))
    num_bins = len(items_in_bins)
    fig, axes = plt.subplots()
    for _, bin_i in enumerate(items_in_bins):
        sum_items = 0
        for item in bin_i:
            axes.bar(_, weights[item], bottom=sum_items, label=f"Item {item}", color=colors(item))
            sum_items += weights[item]
    axes.hlines(max_weight, -0.5, num_bins - 0.5, linestyle="--", color="tab:red", label="Max Weight")
    axes.set_xticks(np.arange(num_bins))
    axes.set_xlabel("Bin")
    axes.set_ylabel("Weight")
    axes.legend()
    if title:
        axes.set_title(title)
    return fig

qaoa = MinimumEigenOptimizer(QAOA(reps=3, quantum_instance=backend))
result_qaoa = qaoa.solve(qubo_simp)
print("Solving QUBO using QAOA")
print(result_qaoa)
print("Solving QUBO using CPLEX")
result_ideal = CplexOptimizer().solve(qubo_simp)
print(result_ideal)

fig = get_figure(interpret(result_ideal.x, weights, num_items, num_bins, simplify=True), weights, max_weight, title="CPLEX")
fig = get_figure(interpret(result_qaoa.x, weights, num_items, num_bins, simplify=True), weights, max_weight, title="QAOA")

max_weight = 15
num_qubits = {}
for i in range(2, 10):
    num_items = i
    num_bins = num_items
    weights = [np.random.randint(1, max_weight) for _ in range(i)]
    num_qubits[i] = BinPacking(num_items, num_bins, weights, max_weight, simplification=True).get_num_vars()

plt.figure()
plt.plot(num_qubits.keys(), num_qubits.values())
plt.xlabel("num items")
plt.ylabel("num qubits")
plt.grid()
print(f"The number of qubits is: {num_qubits}")

x = np.linspace(-2, 2, 100)
plt.figure()
plt.plot(x, np.exp(-x), label="Exponential")
plt.plot(x, 1 - (x) + (x) ** 2 / 2, label="Approximation")
plt.xlabel(r"$t = B - \sum_{i=1}^n s(i)*x_i$", fontsize=14)
plt.ylabel("penalization", fontsize=14)
plt.grid()
plt.legend()

def BinPackingNewApproach(num_items, num_bins, weights, max_weight, alpha=0.01, simplification=False):
    # Construct model using docplex
    mdl = Model("BinPackingNewApproach")

    y = mdl.binary_var_list(num_bins, name="y") # list of variables that represent the bins
    x =  mdl.binary_var_matrix(num_items, num_bins, "x") # variables that represent the items on the specific bin

    objective = mdl.sum(y)
    
    # PENALIZATION
    penalization = 0
    for j in range(num_bins):
        t = max_weight * y[j] - mdl.sum(weights[i] * x[i, j] for i in range(num_items))
        penalization += -t + t**2 / 2
    mdl.minimize(objective + alpha * penalization)

    for i in range(num_items):
        # First set of constraints: the items must be in any bin
        mdl.add_constraint(mdl.sum(x[i, j] for j in range(num_bins)) == 1)

    # Load quadratic program from docplex model
    qp = QuadraticProgram()
    qp.from_docplex(mdl)
    if simplification:
        l = int(np.ceil(np.sum(weights)/max_weight))
        qp = qp.substitute_variables({f"y_{_}":1 for _ in range(l)}) # First simplification 
        qp = qp.substitute_variables({"x_0_0":1}) # Assign the first item into the first bin
        qp = qp.substitute_variables({f"x_0_{_}":0 for _ in range(1, num_bins)}) # as the first item is in the first 
                                                                                #bin it couldn't be in the other bins
    qubo = QuadraticProgramToQubo().convert(qp)# Create a converter from quadratic program to qubo representation
    return qubo

max_weight = 15
num_qubits = {}
num_qubits_new = {}
for i in range(2, 10):
    num_items = i
    num_bins = num_items
    weights = [np.random.randint(1, max_weight) for _ in range(i)]
    num_qubits[i] = BinPacking(num_items, num_bins, weights, max_weight, simplification=True).get_num_vars()
    num_qubits_new[i] = BinPackingNewApproach(num_items, num_bins, weights, max_weight, simplification=True).get_num_vars()
    
plt.figure()
plt.plot(num_qubits.keys(), num_qubits.values(), label="Current Approach")
plt.plot(num_qubits_new.keys(), num_qubits_new.values(), label="New Approach")
plt.xlabel("num items")
plt.ylabel("num qubits")
plt.grid()
plt.legend()
print(f"The number of qubits for the Classical approach is: {num_qubits}")
print(f"The number of qubits is for the new approach is: {num_qubits_new}")



np.random.seed(15)
num_items = 3 # number of items
num_bins = num_items # maximum number of bins
max_weight = 15 # max weight of a bin
weights = np.random.randint(1, max_weight, num_items) # Randomly picking the item weight

qubo_classical = BinPacking(num_items, num_bins, weights, max_weight, simplification=True)
qubo_new = BinPackingNewApproach(num_items, num_bins, weights, max_weight, alpha=0.05, simplification=True)
qaoa = MinimumEigenOptimizer(QAOA(optimizer=COBYLA(maxiter=10), reps=3, quantum_instance=backend))
result_qaoa_classical = qaoa.solve(qubo_classical)
result_qaoa_new = qaoa.solve(qubo_new)
result_cplex_classical = CplexOptimizer().solve(qubo_classical)

print("------------------------------------------")
print("Solving QUBO using QAOA Classical Approach")
print(result_qaoa_classical)
print("------------------------------------------")
print("Solving QUBO using QAOA New Approach")
print(result_qaoa_new)
print("------------------------------------------")
print("Solving QUBO using CPLEX")
print(result_cplex_classical)
print("------------------------------------------")


fig = get_figure(interpret(result_ideal_old.x, weights, num_items, num_bins, simplify=True), weights, max_weight, title="CPLEX")
fig = get_figure(interpret(result_qaoa_old.x, weights, num_items, num_bins, simplify=True), weights, max_weight, title="QAOA classical approach")
fig = get_figure(interpret(result_qaoa_new.x, weights, num_items, num_bins, simplify=True), weights, max_weight, title="QAOA New Approach")

solutions_new = []
solutions_classical = []
for _ in range(10):
    result_qaoa_new = qaoa.solve(qubo_new)
    result_qaoa_classical = qaoa.solve(qubo_classical)
    solutions_new.append(result_qaoa_new.fval)
    solutions_classical.append(result_qaoa_classical.fval)

num_vars = qubo_new.get_num_vars()
vars_new = qubo_new.variables
result_cplex = CplexOptimizer().solve(qubo_classical)
result_new_ideal = qubo_new.objective.evaluate(result_cplex.x[:num_vars])# Replacing the ideal solution into
                                               #our new approach to see the optimal solution on the new objective
                                               #function

ratio_new = (np.abs(1 - np.array(solutions_new) / result_new_ideal) < 0.05).sum()
ratio_classical = (np.abs(1 - np.array(solutions_classical) / result_cplex.fval) < 0.05).sum()

fig, ax = plt.subplots(1,2, figsize=(15,5))
ax[0].plot(range(1,11), solutions_new, color="tab:green",label="new QUBO QAOA")
ax[0].plot(range(1,11), solutions_classical, color="tab:red",label="classical QUBO QAOA")
ax[0].plot([1, 10], 2*[result_new_ideal], ":", color="tab:green", label="new QUBO optimal")
ax[0].plot([1, 10], 2*[result_cplex.fval], ":", color="tab:red", label="classical QUBO optimal")
ax[0].set_xlabel("n", fontsize=14)
ax[0].set_ylabel("Cost Obj. function", fontsize=14)
ax[0].grid()
ax[0].legend()

ax[1].bar(0, 100 * ratio_new / 10, label="new")
ax[1].bar(1, 100 * ratio_classical/ 10 + 0.5, label="Classical")
ax[1].legend()
ax[1].set_xticks([0,1])
ax[1].set_xticklabels(["", ""])
ax[1].set_ylabel("Ratio", fontsize=14)

def new_eq_optimal(qubo_new, qubo_classical):
    """
    From the classical solution and considering that cplex solution is the optimal, we can traslate the optimal
    solution to the QUBO representation based on our approach.
    
    
    """
    num_vars = qubo_new.get_num_vars()
    vars_new = qubo_new.variables
    result_cplex = CplexOptimizer().solve(qubo_classical)
    result_new_ideal = qubo_new.objective.evaluate(result_cplex.x[:num_vars])# Replacing the ideal solution into
                                               #our new approach to see the optimal solution on the new objective
                                               #function
    return result_new_ideal

np.random.seed(1)
num_items = 3 # number of items
num_bins = num_items # maximum number of bins
max_weight = 15 # max weight of a bin
repetitions = 10
cases = 5
solutions_new = {}
optimal_new = []
ratio_new = []
for i in range(cases): # Testing 5 different randomly selected configurations of the problem
    print(f"----------- Case {i+1} -------------")
    weights = np.random.randint(1, max_weight, num_items) # Randomly picking the item weight
    qubo_new = BinPackingNewApproach(num_items, num_bins, weights, max_weight, alpha=0.05, simplification=True)
    qubo_classical = BinPacking(num_items, num_bins, weights, max_weight, simplification=True)
    optimal_new.append(new_eq_optimal(qubo_new, qubo_classical))
    solutions_new[i] = []
    for _ in range(repetitions): #Testing the ratio
        result_qaoa_new = qaoa.solve(qubo_new)
        solutions_new[i].append(result_qaoa_new.fval)
    ratio_new.append(100 * ((np.array(solutions_new[i]) / optimal_new[-1] - 1) < 5e-2).sum() / repetitions)

plt.figure()
plt.bar([i for i in range(1, cases + 1)], ratio_new)
plt.xlabel("Case")
plt.ylabel("Ratio")

np.random.seed(1)
num_items = 4 # number of items
num_bins = num_items # maximum number of bins
max_weight = 15 # max weight of a bin
repetitions = 10
cases = 5
solutions_new = {}
optimal_new = []
ratio_new = []
for i in range(cases): # Testing 5 different randomly selected configurations of the problem
    print(f"----------- Case {i+1} -------------")
    weights = np.random.randint(1, max_weight, num_items) # Randomly picking the item weight
    qubo_new = BinPackingNewApproach(num_items, num_bins, weights, max_weight, alpha=0.05, simplification=True)
    qubo_classical = BinPacking(num_items, num_bins, weights, max_weight, simplification=True)
    optimal_new.append(new_eq_optimal(qubo_new, qubo_classical))
    solutions_new[i] = []
    for _ in range(repetitions): #Testing the ratio
        result_qaoa_new = qaoa.solve(qubo_new)
        solutions_new[i].append(result_qaoa_new.fval)
    ratio_new.append(100 * ((np.array(solutions_new[i]) / optimal_new[-1] - 1) < 5e-2).sum() / repetitions)

plt.figure()
plt.bar([i for i in range(1, cases + 1)], ratio_new)
plt.xlabel("Case")
plt.ylabel("Ratio")

np.random.seed(1)
num_items = 5 # number of items
num_bins = num_items # maximum number of bins
max_weight = 15 # max weight of a bin
repetitions = 10
cases = 5
solutions_new = {}
optimal_new = []
ratio_new = []
qaoa = MinimumEigenOptimizer(QAOA(optimizer=COBYLA(maxiter=100), reps=4, quantum_instance=backend))
for i in range(cases): # Testing 5 different randomly selected configurations of the problem
    print(f"----------- Case {i+1} -------------")
    weights = np.random.randint(1, max_weight, num_items) # Randomly picking the item weight
    qubo_new = BinPackingNewApproach(num_items, num_bins, weights, max_weight, alpha=0.05, simplification=True)
    qubo_classical = BinPacking(num_items, num_bins, weights, max_weight, simplification=True)
    optimal_new.append(new_eq_optimal(qubo_new, qubo_classical))
    solutions_new[i] = []
    for _ in range(repetitions): #Testing the ratio
        result_qaoa_new = qaoa.solve(qubo_new)
        solutions_new[i].append(result_qaoa_new.fval)
    ratio_new.append(100 * ((np.array(solutions_new[i]) / optimal_new[-1] - 1) < 5e-2).sum() / repetitions)

plt.figure()
plt.bar([i for i in range(1, cases + 1)], ratio_new)
plt.xlabel("Case")
plt.ylabel("Ratio")

