import numpy as np

from qiskit import IBMQ, Aer
from qiskit.algorithms import QAOA, VQE, NumPyMinimumEigensolver
from qiskit.algorithms.optimizers import COBYLA
from qiskit.utils import QuantumInstance, algorithm_globals
from qiskit.providers.aer.noise.noise_model import NoiseModel

from qiskit_optimization import QuadraticProgram
from qiskit_optimization.algorithms import MinimumEigenOptimizer
from qiskit_optimization.converters import QuadraticProgramToQubo

import qiskit.test.mock as Fake

IBMQ.load_account()
IBMQ.providers()

algorithm_globals.massive = True

quadprog = QuadraticProgram("test")
quadprog.integer_var(name="x_1", lowerbound=0, upperbound=4)
quadprog.integer_var(name="x_2", lowerbound=-2, upperbound=2)
quadprog.integer_var(name="x_3", lowerbound=-2, upperbound=4)
quadprog.minimize(
    linear={"x_3": -6},
    quadratic={("x_1", "x_1"): 1, ("x_2", "x_2"): 1, ("x_1", "x_2"): -1}
)
quadprog.linear_constraint(linear={"x_1": 1, "x_2": 1}, sense="=", rhs=2)
quadprog.quadratic_constraint(quadratic={("x_2", "x_3"): 1}, sense=">=", rhs=1)
print(quadprog.export_as_lp_string())

def cropyield_quadratic_program(n=3):
    cropyield = QuadraticProgram()
    
    # initialize the crop types as integer variables
    crops = ["Wheat", "Soybeans", "Maize", "PushPull"]
    w = "Wheat"
    s = "Soybeans"
    m = "Maize"
    p = "PushPull"
    for crop in crops:
        cropyield.integer_var(name=crop, lowerbound=0, upperbound=n)
        
    l = {}
    q = {}
    # initialize the equation for total crop output. 
    # these represent how "good" each of the crop types are in terms of the final equation
    # so wheat is 2x as good as soybeans and maize is 4x as good, and pushpull provides no direct benefit
    l[w] = 2
    l[s] = 1
    l[m] = 4
    
    # define the intercropping coefficients, aka how well certain crops will be when grown together
    q[(w, s)] = 2.4
    q[(w, m)] = 4
    q[(w, p)] = 4
    q[(s, m)] = 2
    q[(s, p)] = 1
    q[(m, p)] = 5
    
    # we want to maximize the output
    cropyield.maximize(linear=l, quadratic=q)
    
    # but make sure that the sum of the number of squares is <= 3
    cropyield.linear_constraint(linear={w:1, s:1, m:1, p:1}, sense="<=", rhs=n)
    
    return cropyield

# convert the quadratic problem to a quantum problem
cropyield = cropyield_quadratic_program(n=3)
ising_operations, _ = (
    QuadraticProgramToQubo()
    .convert(
        cropyield,
    )
    .to_ising()
)
print(f"Number of qubits required is {ising_operations.num_qubits}")

QuadraticProgramToQubo().convert(cropyield)

# use the Aer simulator
backend = Aer.get_backend("qasm_simulator")

algorithm_globals.random_seed = 271828

def get_classical_solution_for(quadprog: QuadraticProgram):
    solver = NumPyMinimumEigensolver()
    optimizer = MinimumEigenOptimizer(solver)
    return optimizer.solve(quadprog)

def get_QAOA_solution_for(
    quadprog: QuadraticProgram, quantumInstance: QuantumInstance, optimizer=None
):
    _eval_count = 0
    
    def callback(eval_count, parameters, mean, std):
        nonlocal _eval_count
        _eval_count = eval_count
    
    solver = QAOA(
        optimizer=optimizer, quantum_instance=quantumInstance, callback=callback
    )
    
    optimizer = MinimumEigenOptimizer(solver)
    result = optimizer.solve(quadprog)
    return result, _eval_count

def get_VQE_solution_for(
    quadprog: QuadraticProgram, quantumInstance: QuantumInstance, optimizer=None,
):
    _eval_count = 0

    def callback(eval_count, parameters, mean, std):
        nonlocal _eval_count
        _eval_count = eval_count

    # Create solver and optimizer
    solver = VQE(
        optimizer=optimizer, quantum_instance=quantumInstance, callback=callback
    )

    # Create optimizer for solver
    optimizer = MinimumEigenOptimizer(solver)

    # Get result from optimizer
    result = optimizer.solve(quadprog)
    return result, _eval_count

classical_result = get_classical_solution_for(cropyield)

print("Solution found using classical method:\n")
print(f"Maximum crop yield is {classical_result.fval} tons")
print("Crops used are: ")

_crops = [v.name for v in cropyield.variables]
for cropIndex, cropHectares in enumerate(classical_result.x):
    print(f"\t{cropHectares} ha of {_crops[cropIndex]}")
    

simulator_instance = QuantumInstance(
    backend=backend,
    seed_simulator=algorithm_globals.random_seed,
    seed_transpiler=algorithm_globals.random_seed,
)

qaoa_result, qaoa_eval_count = get_QAOA_solution_for(cropyield, simulator_instance)

print("Solution found using the QAOA method:\n")
print(f"Maximum crop-yield is {qaoa_result.fval} tons")
print(f"Crops used are: ")
for cropHectares, cropName in zip(qaoa_result.x, qaoa_result.variable_names):
    print(f"\t{cropHectares} ha of {cropName}")

print(f"\nThe solution was found within {qaoa_eval_count} evaluations of QAOA.")

# Create a QuantumInstance
simulator_instance = QuantumInstance(
    backend=backend,
    seed_simulator=algorithm_globals.random_seed,
    seed_transpiler=algorithm_globals.random_seed,
)

# Get VQE result
vqe_result, vqe_eval_count = get_VQE_solution_for(
    cropyield, simulator_instance)

# Format and print result
print("Solution found using the VQE method:\n")
print(f"Maximum crop-yield is {vqe_result.fval} tons")
print(f"Crops used are: ")
for cropHectares, cropName in zip(vqe_result.x, vqe_result.variable_names):
    print(f"\t{cropHectares} ha of {cropName}")

print(f"\nThe solution was found within {vqe_eval_count} evaluations of VQE")


problem = cropyield_quadratic_program(n=50)

for _backend in IBMQ.get_provider(hub='ibm-q', group='open', project='main').backends():
    print(_backend.name())

provider = IBMQ.get_provider(hub='ibm-q', group='open', project='main')
backend_real = provider.get_backend('ibmq_qasm_simulator')
quantum_instance_real = QuantumInstance(backend_real, shots=2048)
optimizer = COBYLA(maxiter=1)


## Get result from real device with VQE
vqe_result_real, vqe_eval_count_real = get_VQE_solution_for(
    problem, quantum_instance_real, optimizer=optimizer
)


# Format and print result
print("Solution found using the VQE method:\n")
print(f"Maximum crop-yield is {vqe_result_real.fval} tons")
print(f"Crops used are: ")
for cropHectares, cropName in zip(vqe_result_real.x, vqe_result_real.variable_names):
    print(f"\t{cropHectares} ha of {cropName}")

print(
    f"\nThe solution was found within {vqe_eval_count_real} evaluations of VQE")


