from math import pi
import numpy as np
from qiskit_nature.operators.second_quantization.fermionic_op import FermionicOp
from qiskit_nature.problems.second_quantization.lattice import HyperCubicLattice, LineLattice, SquareLattice, TriangularLattice
from qiskit_nature.problems.second_quantization.lattice import FermiHubbardModel

num_nodes = 6
boundary_condition = 'open'
line = LineLattice(
    num_nodes = num_nodes,
    boundary_condition=boundary_condition
)

line.draw()

num_nodes = 11
boundary_condition = 'periodic' # `boundary_condition` must be 'open' or 'periodic'.
line_p = LineLattice(
    num_nodes = num_nodes,
    boundary_condition=boundary_condition
)

line_p.draw()

rows = 4
cols = 3
boundary_condition = 'open'
square = SquareLattice(
    rows=rows,
    cols=cols,
    boundary_condition=boundary_condition
)

square.draw()

rows = 4
cols = 3
boundary_condition = ('open', 'periodic') # open in the x-direction, periodic in the y-direction
square_op = SquareLattice(
    rows=rows,
    cols=cols,
    boundary_condition=boundary_condition
)

square_op.draw()

size = (3, 4, 5)
boundary_condition = ('open', 'open', 'open')
hypercubic = HyperCubicLattice(
    size=size,
    boundary_condition=boundary_condition
)

# function for setting the positions
def indextocoord_3d(index:int, size:tuple, angle) -> list:
    z = index // (size[0]*size[1])
    a = index % (size[0]*size[1])
    y = a // size[0]
    x = a % size[0]
    vec_x = np.array([1, 0])
    vec_y = np.array([np.cos(angle), np.sin(angle)])
    vec_z = np.array([0, 1])
    return_coord = x*vec_x + y*vec_y + z*vec_z
    return return_coord.tolist()

pos = dict([(index, indextocoord_3d(index, size, angle=pi/4)) for index in range(np.prod(size))])

hypercubic.draw(pos=pos)

rows = 4
cols = 3
boundary_condition = 'open'

triangular = TriangularLattice(
    rows=rows,
    cols=cols,
    boundary_condition=boundary_condition
)

triangular.draw()

line = LineLattice(
    num_nodes=3,
    boundary_condition='open'
)

line.draw()

t = -1.0
t_onsite = 0.0
U = 5.0

fhm = FermiHubbardModel.uniform_parameters(
    lattice=line,
    uniform_hopping=t, 
    uniform_onsite_potential=t_onsite, 
    onsite_interaction=U)

FermionicOp.set_truncation(0)
print(fhm.second_q_ops(display_format='sparse'))

from qiskit_nature.converters.second_quantization import QubitConverter
from qiskit_nature.mappers.second_quantization import JordanWignerMapper
from qiskit.algorithms import VQE
from qiskit.algorithms.optimizers import SLSQP
from qiskit.circuit.library import RealAmplitudes
from qiskit.opflow import AerPauliExpectation
from qiskit.providers.aer import AerSimulator

# make a lattice
line = LineLattice(
    num_nodes=3,
    boundary_condition='open'
)

# the FHM
t = -1.0
t_onsite = 0.0
U = 5.0

fhm = FermiHubbardModel.uniform_parameters(
    lattice=line,
    uniform_hopping=t,
    uniform_onsite_potential=t_onsite,
    onsite_interaction=U
)

# the FHM in terms of `FermionicOp`
ham = fhm.second_q_ops(display_format='sparse')
print(ham.reduce())

# convert the Hamiltonian to qubits
qubit_converter = QubitConverter(mapper=JordanWignerMapper())
qubit_op = qubit_converter.convert(ham)

# define the quantum circuit
circ = RealAmplitudes(num_qubits=ham.register_length, reps=2)
print('Quantum circuit')
print(circ.decompose().draw())

# make the vqe instance
vqe = VQE(
    ansatz=circ,
    optimizer=SLSQP(),
    expectation=AerPauliExpectation(),
    max_evals_grouped=9999,
    quantum_instance=AerSimulator()
            )

result = vqe.compute_minimum_eigenvalue(operator=qubit_op)
energy = np.real(result.eigenvalue)

print(f'Energy: {energy}')

