import matplotlib.pyplot as plt
import networkx as nx
import pandas as pd
import pyomo.environ as pyo
from IPython.display import Markdown, display
from pyomo.environ import value

# Dimensions of the container (width and height)
CONTAINER_WIDTH = 4
CONTAINER_HEIGHT = 2

# List of rectangles with (width, height) tuples
RECTANGLES = [(1, 1), (2, 2), (2, 1)]

from pyomo.environ import (
    Binary,
    ConcreteModel,
    Constraint,
    Objective,
    RangeSet,
    SolverFactory,
    Var,
)


def define_rectangkes_packing_model(rectangles, container_width, container_height):
    # Number of rectangles
    num_rectangles = len(rectangles)

    # Create a model
    model = ConcreteModel()

    # Sets for rectangles and grid positions
    model.R = RangeSet(0, num_rectangles - 1)
    model.W = RangeSet(0, container_width - 1)
    model.H = RangeSet(0, container_height - 1)

    # Binary variable: 1 if rectangle r is placed at position (i, j), 0 otherwise
    model.place = Var(model.R, model.W, model.H, domain=Binary)

    # Constraints to ensure each rectangle is placed at most once
    def one_place_rule(model, r):
        return sum(model.place[r, i, j] for i in model.W for j in model.H) <= 1

    model.one_place = Constraint(model.R, rule=one_place_rule)

    # Constraints to ensure rectangles do not overlap -  we ensure for each coordinate that it is occupied by at most 1 rectangle
    def non_overlap_rule(model, i, j):
        return (
            sum(
                model.place[r, i2, j2]
                for r in model.R
                for i2 in range(i - rectangles[r][0] + 1, i + 1)
                for j2 in range(j - rectangles[r][1] + 1, j + 1)
                if i2 >= 0 and j2 >= 0
            )
            <= 1
        )

    model.non_overlap = Constraint(model.W, model.H, rule=non_overlap_rule)

    # Constraints to ensure each rectangle is within the container
    def within_container_rule(model, r, i, j):
        if (
            i + rectangles[r][0] > container_width
            or j + rectangles[r][1] > container_height
        ):
            return model.place[r, i, j] == 0
        return Constraint.

Skip

    model.within_container = Constraint(
        model.R, model.W, model.H, rule=within_container_rule
    )

    # Objective function: maximize the number of placed rectangles
    model.obj = Objective(
        expr=-sum(
            model.place[r, i, j] for r in model.R for i in model.W for j in model.H
        ),
        sense=pyo.minimize,
    )

    return model

model = define_rectangkes_packing_model(RECTANGLES, CONTAINER_WIDTH, CONTAINER_HEIGHT)

model.pprint()

from classiq import *
from classiq.applications.combinatorial_optimization import OptimizerConfig, QAOAConfig

qaoa_config = QAOAConfig(num_layers=10, penalty_energy=100)

optimizer_config = OptimizerConfig(max_iteration=60, alpha_cvar=1)

qmod = construct_combinatorial_optimization_model(
    pyo_model=model,
    qaoa_config=qaoa_config,
    optimizer_config=optimizer_config,
)

from classiq.execution import ClassiqBackendPreferences

qmod = set_execution_preferences(
    qmod,
    num_shots=10000,
    backend_preferences=ClassiqBackendPreferences(backend_name="simulator"),
)

qprog = synthesize(qmod)
show(qprog)

result = execute(qprog).result_value()

result.convergence_graph

import pandas as pd

from classiq.applications.combinatorial_optimization import (
    get_optimization_solution_from_pyo,
)

solution = get_optimization_solution_from_pyo(
    model, vqe_result=result, penalty_energy=qaoa_config.penalty_energy
)
optimization_result = pd.DataFrame.from_records(solution)
optimization_result.sort_values(by="cost", ascending=True).head(5)

optimization_result.hist("cost", weights=optimization_result["probability"])

best_solution = optimization_result.solution[optimization_result.cost.idxmin()]

def extract_no_boundries_qubit_indexes(model):
    no_qubits_indexes = []
    for c in model.component_objects(Constraint, active=True):
        cdata = getattr(model, c.name)
        for index in cdata:
            lb = value(cdata[index].lower)
            ub = value(cdata[index].upper)
            if lb == ub:
                no_qubits_indexes.append(index)
    return no_qubits_indexes

def prepare_solution_for_floorplan_visual(best_solution, no_qubits_indexes):
    best_solution_index_count = 0
    solution = {}
    for r in model.R:
        for w in model.W:
            for h in model.H:
                if (r, w, h) not in no_qubits_indexes:
                    solution[(r, w, h)] = best_solution[best_solution_index_count]
                    best_solution_index_count += 1
                else:
                    solution[(r, w, h)] = 0
    return solution


def place_rectangles(solution):
    placement = {}
    for r in model.R:
        for i in model.W:
            for j in model.H:
                if solution[(r, i, j)] == 1:
                    placement[r] = (i, j)
    return placement

import matplotlib.patches as patches
import matplotlib.pyplot as plt


# Function to visualize the rectangle packing solution
def visualize_packing(container_width, container_height, rectangles, placement):
    fig, ax = plt.subplots(1)
    ax.set_xlim(0, container_width)
    ax.set_ylim(0, container_height)

    # Draw the container
    container = patches.Rectangle(
        (0, 0),
        container_width,
        container_height,
        linewidth=1,
        edgecolor="r",
        facecolor="none",
    )
    ax.add_patch(container)

    # Draw each rectangle
    colors = [
        "blue",
        "green",
        "orange",
        "purple",
        "yellow",
    ]  # Add more colors if needed
    for r, (i, j) in placement.items():
        width, height = rectangles[r]
        rect = patches.Rectangle(
            (i, j),
            width,
            height,
            linewidth=1,
            edgecolor="black",
            facecolor=colors[r % len(colors)],
            alpha=0.5,
        )
        ax.add_patch(rect)
        plt.text(
            i + width / 2,
            j + height / 2,
            f"{r}",
            ha="center",
            va="center",
            color="white",
        )

    plt.gca().set_aspect("equal", adjustable="box")
    plt.gca().invert_yaxis()  # Invert y axis to match the typical matrix/grid representation
    # plt.grid(True)
    plt.show()


# Visualize the solution
visualize_packing(
    CONTAINER_WIDTH,
    CONTAINER_HEIGHT,
    RECTANGLES,
    place_rectangles(
        prepare_solution_for_floorplan_visual(
            best_solution, extract_no_boundries_qubit_indexes(model)
        )
    ),
)

"""
# Create a solver
solver = SolverFactory('glpk')


# Solve the model
solver.solve(model)

# Display results
for r in model.R:
    placed = [(i, j) for i in model.W for j in model.H if model.place[r, i, j].value == 1]
    if placed:
        print(f"Rectangle {RECTANGLES[r]} placed at position {placed[0]}")
    else:
        print(f"Rectangle {r} not placed")

import matplotlib.pyplot as plt
import matplotlib.patches as patches


placement = {}
for r in model.R:
    for i in model.W:
        for j in model.H:
            if model.place[r, i, j].value == 1:
                placement[r] = (i, j)
                


# Visualize the solution
visualize_packing(CONTAINER_WIDTH,CONTAINER_HEIGHT, RECTANGLES, placement)
"""