# If using this on Google colab, we need to install the packages
try:
  import google.colab
  IN_COLAB = True
except:
  IN_COLAB = False

if IN_COLAB:
    !pip install eqc_models pyomo

# Import the QCI models and solvers
import eqc_models
from eqc_models.base import ConstrainedPolynomialModel, PolynomialModel, QuadraticModel
from eqc_models.solvers import Dirac3IntegerCloudSolver, Dirac3ContinuousCloudSolver

# Import Matplotlib to generate plots
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker

# Import numpy and scipy for certain numerical calculations below
import numpy as np
from scipy.special import gamma
import pyomo.environ as pyo
import os

if IN_COLAB:
    !pip install idaes-pse --pre
    !idaes get-extensions --to ./bin 
    os.environ['PATH'] += ':bin'

# Define the API URL and token for QCI
# Note: Replace the api_url and api_token with your actual QCI API URL and token
api_url = "https://api.qci-prod.com" 
api_token = "" # Replace with your actual API token

# 1. Create a concrete Pyomo model
m = pyo.ConcreteModel(name="Simple_Quadratic_Program")

# 2. Define the decision variables with their bounds
# The variables x_1, x_2, and x_3 are defined with bounds [0, 10].
m.x = pyo.Var([1, 2, 3], domain=pyo.NonNegativeReals, bounds=(0, 10))

# 3. Define the objective function
# The objective is to minimize the quadratic expression.
m.obj = pyo.Objective(
    expr=3*m.x[1]**2 + 2*m.x[2]**2 + m.x[3]**2, 
    sense=pyo.minimize
)

# 4. Define the linear constraint
# The sum of the variables must equal 10.
m.c1 = pyo.Constraint(expr=m.x[1] + m.x[2] + m.x[3] == 10)

# 5. Create a solver instance and solve the model
# We specify 'ipopt' as the solver. 'tee=True' displays the solver's log.
solver = pyo.SolverFactory('ipopt')
results = solver.solve(m, tee=False)

# 6. Display the optimization results
print("\n" + "="*30)
print("-- 📊 Optimization Results --")
print(f"Solver Status: {results.solver.status}")
print(f"Termination Condition: {results.solver.termination_condition}")
print(f"Objective Value (E): {pyo.value(m.obj):.4f}")
print("\n-- 🏗️ Variable Values --")
for i in m.x:
    print(f"x[{i}] = {pyo.value(m.x[i]):.4f}")
print("="*30)

R = 10
coefficients = [3, 2, 1]
indices = [(1,1), (2,2), (3,3)] # Create a polynomial model
model = PolynomialModel(coefficients, indices)
model.upper_bound = R*np.ones((3,)) # Bounds for the variables

solver = Dirac3ContinuousCloudSolver(url=api_url, api_token=api_token)
response = solver.solve(model, sum_constraint=10, relaxation_schedule=1, num_samples=20) # sum constraint is set to 10, 5 samples are collected

# Print the results
print(response["results"]["solutions"])
print(response["results"]["energies"])
print(response["results"]["counts"])

print(response)

# Extract the results from the response
results = response["results"]
energies = results["energies"]
solutions = results["solutions"]

# Find the index of the best solution (minimum energy)
best_index = np.argmin(energies)

# Get the best solution and its energy
best_energy = energies[best_index]
best_solution = solutions[best_index]

# Print the results in a clear format
print("-- 📊 Best Solution Found --")
print(f"Optimal Objective (Energy): {best_energy:.4f}")
print(f"Variable Values: {np.round(best_solution, 4)}")

# 1. Create a concrete Pyomo model
m = pyo.ConcreteModel(name="Simple_QUBO")

# 2. Define the binary variables
# x[1] and x[2] can only take values of 0 or 1.
m.x = pyo.Var([1, 2], domain=pyo.Binary)

# 3. Define the objective function
# We use the property that for a binary variable x, x^2 = x.
# So, the objective E = -0.75*x1^2 - 0.25*x2^2 + 2*x1*x2 + 0.5*x1 becomes:
# E = -0.75*x1 - 0.25*x2 + 2*x1*x2
m.obj = pyo.Objective(
    expr= -0.75 * m.x[1] - 0.25 * m.x[2] + 2 * m.x[1] * m.x[2] + 0.5 * m.x[1],
    sense=pyo.minimize
)

# 4. Create a solver instance and solve the model
# This is an unconstrained problem, so no constraints are needed.
# We use a Mixed-Integer Programming (MINLP) solver like 'bonmin'.
solver = pyo.SolverFactory('bonmin')
results = solver.solve(m, tee=False)

# 5. Display the optimization results
print("\n" + "="*30)
print("-- 📊 Optimization Results --")
print(f"Solver Status: {results.solver.status}")
print(f"Termination Condition: {results.solver.termination_condition}")
print(f"Objective Value: {pyo.value(m.obj):.4f}")
print("\n-- 🏗️ Variable Values --")
for i in m.x:
    print(f"x[{i}] = {int(pyo.value(m.x[i]))}")
print("="*30)

coefficients = [-0.75, -0.25, 2, 0.5]  # Coefficients for the polynomial model
indices = [(1, 1), (2, 2), (1, 2), (0, 1)]  # Create a polynomial model
model = PolynomialModel(coefficients, indices)

# add upper bounds for the variables
model.upper_bound = np.ones((2,))  # Bounds for the variables

solver = Dirac3IntegerCloudSolver(url=api_url, api_token=api_token)
response = solver.solve(model, relaxation_schedule=1, num_samples=100)

# Print the results
print(response)
print(response["results"]["solutions"])
print(response["results"]["energies"])
print(response["results"]["counts"])

def plot_solution_distribution(response):
    """
    Processes a solver response and plots the distribution of unique solutions.

    Args:
        response (dict): The response dictionary from the Dirac-3 solver,
                         containing 'solutions' and 'counts' in its 'results' key.
    """
    # --- Data Processing ---
    # Extract the raw results from the response dictionary
    results = response.get('results', {})
    raw_solutions = results.get('solutions', [])
    raw_counts = results.get('counts', [])

    if not raw_solutions or not raw_counts:
        print("Response object does not contain valid solutions or counts.")
        return

    # Aggregate counts for each unique solution
    solution_counts = {}
    for solution, count in zip(raw_solutions, raw_counts):
        # Convert list to tuple to use it as a dictionary key
        solution_tuple = tuple(solution)
        solution_counts[solution_tuple] = solution_counts.get(solution_tuple, 0) + count

    # Prepare data for plotting
    # Create string labels for the x-axis from the solution tuples
    plot_solutions_str = [str(list(s)) for s in solution_counts.keys()]
    plot_counts = list(solution_counts.values())

    # --- Plotting ---
    plt.style.use('seaborn-v0_8-whitegrid')
    fig, ax = plt.subplots(figsize=(8, 7)) # Increased height for rotated labels

    # Create the bar plot
    bars = ax.bar(plot_solutions_str, plot_counts, color='#4A90E2', width=0.5)

    # Add labels and title for clarity
    ax.set_xlabel("Solution", fontsize=12, labelpad=15)
    ax.set_ylabel("Counts (Frequency)", fontsize=12)
    ax.set_title("Distribution of Solutions", fontsize=14, fontweight='bold')
    
    # Rotate x-axis labels to prevent overlap
    plt.xticks(rotation=90)
    
    # Ensure the y-axis only displays integer values
    ax.yaxis.set_major_locator(mticker.MaxNLocator(integer=True))

    # Add count labels on top of each bar
    for bar in bars:
        yval = bar.get_height()
        ax.text(bar.get_x() + bar.get_width()/2.0, yval + 0.1, int(yval), ha='center', va='bottom', fontsize=11)

    # Set y-axis to start at 0 and give some space at the top
    if plot_counts:
        ax.set_ylim(0, max(plot_counts) * 1.15)

    # Adjust layout to make room for rotated labels
    plt.tight_layout()
    plt.show()

# Plot the distribution of unique solutions
plot_solution_distribution(response)

A = np.array([[1, 0, 0, 1, 1, 1, 0, 1, 1, 1, 1],
            [0, 1, 0, 1, 0, 1, 1, 0, 1, 1, 1],
            [0, 0, 1, 0, 1, 0, 1, 1, 1, 1, 1]])
b = np.array([1, 1, 1])
c = np.array([2, 4, 4, 4, 4, 4, 5, 4, 5, 6, 5])
num_variables = A.shape[1]

# Build a Pyomo model for the constrained linear integer program
m = pyo.ConcreteModel(name="Constrained_Linear_Integer_Program")

# Define the set of variable indices
m.I = pyo.RangeSet(0, num_variables - 1)
# Define the set of constraint indices
m.J = pyo.RangeSet(0, A.shape[0] - 1)

# Define the 11 binary variables
m.x = pyo.Var(m.I, domain=pyo.Binary)

# Define the Objective Function
# The objective is the original linear expression.
m.obj = pyo.Objective(
    expr=sum(c[i] * m.x[i] for i in m.I),
    sense=pyo.minimize
)

# Define the constraints
# We define a set of constraints Ax = b directly.
def ax_constraint_rule(model, j):
    # For each row j, sum(A[j,i] * x[i]) must equal b[j]
    return sum(A[j, i] * model.x[i] for i in model.I) == b[j]

m.constraints = pyo.Constraint(m.J, rule=ax_constraint_rule)

# Solve the model
# We use a Mixed-Integer Linear Programming (MILP) solver.
solver = pyo.SolverFactory('cbc')
results = solver.solve(m, tee=False)

# Display the optimization results
print("\n" + "="*40)
print("-- 📊 Optimization Results --")
print(f"Solver Status: {results.solver.status}")
print(f"Termination Condition: {results.solver.termination_condition}")
print(f"Objective Value: {pyo.value(m.obj):.4f}")

print("\n-- 🏗️ Variable Values --")
solution_vector = [int(pyo.value(m.x[i])) for i in m.I]
print(f"x = {solution_vector}")

epsilon = 1 # small value to avoid division by zero
rho = np.sum(np.abs(c)) + epsilon # penalty multiplier

# Construct the quadratic Hamiltonian J
J = rho*np.matmul(A.T,A)

# Construct the linear Hamiltonian C
C = c.copy()  
C -= rho*2*(np.matmul(b.T,A))

# Construct the offset term cQ
cQ = rho*np.matmul(b.T,b)

# --- Formatted Printing ---
print("\n" + "="*60)
print("--- QUBO Formulation Parameters ---")
# Set numpy print options for better readability
np.set_printoptions(precision=2, suppress=True)

print("\nQuadratic Hamiltonian (J):\n", J)
print("\n" + "-"*60)
print("\nLinear Hamiltonian (C):\n", C)
print("\n" + "-"*60)
print(f"\nOffset Term (cQ): {cQ:.2f}")
print("\n" + "="*60)

model = QuadraticModel(C, J)
model.upper_bound = np.ones((11, )) # Set the upper bound for the variables

solver = Dirac3IntegerCloudSolver(url=api_url, api_token=api_token)
Q = model.qubo.Q

response = solver.solve(model, num_samples=10) # qubomodel

solution = model.decode(np.array(response["results"]["solutions"][0]), "qubo")

true_solution = model.evaluate(solution) + cQ # add the offset term to get the true objective value

print(f"Decoded Solution (x): {solution}")
print(f"Objective Value (E): {true_solution}")

# Plot the distribution of unique solutions
plot_solution_distribution(response)

coefficients = c
num_vars = 11
indices = [[0, i + 1] for i in range(num_vars)]

lhs = A
rhs = b

# Create the constrained polynomial model
constraint_model = ConstrainedPolynomialModel(coefficients, indices, lhs, rhs)
constraint_model.upper_bound = np.ones(num_vars, ) # Set the upper bound for the variables
# Set the penalty multiplier for the constraints
# This is a crucial step to ensure the constraints are satisfied in the optimization process.
constraint_model.penalty_multiplier = np.sum(np.abs(c)) + 1  # Set the penalty multiplier
print(constraint_model.offset) # Offse term generated by the model

solver = Dirac3IntegerCloudSolver(url=api_url, api_token=api_token)
response = solver.solve(model, num_samples=10)

# plot the distribution of unique solutions
plot_solution_distribution(response)

# Extract the best solution and its objective value
best_solution = response["results"]["solutions"][0]
best_objective_value = response["results"]["energies"][0] + constraint_model.offset * constraint_model.penalty_multiplier 
# Add the offset term to get the true objective value

print(f"Best Solution: {best_solution}")
print(f"Best Objective Value: {best_objective_value}")