%load_ext autoreload
%autoreload 2

from qiskit_cold_atom.applications import FermiHubbard1D

# defining the system 
system = FermiHubbard1D(
    num_sites = 3,          # number of lattice sites
    particles_up = 1,       # number of spin up particles
    particles_down = 1,     # number of spin down particles
    hop_strength = 2.,      # parameter J tuning the hopping 
    int_strength = 1,       # parameter U tuning the interaction
    potential = [1, 0, -1]  # parameters mu tuning the local potential
)

from qiskit_cold_atom.fermions.fermion_circuit_solver import FermionicState

# defining the initial state with one spin_up particle in the left site
# and one spin-down particle in the right site
initial_state = FermionicState([[1, 0, 0], [0, 0, 1]])

from qiskit_nature.second_q.operators import FermionicOp

# Observables: Spin-density and particle density at the first site
spin_density = FermionicOp({'+_0 -_0': 1, '+_3 -_3': -1}, num_spin_orbitals=6)
particle_density = FermionicOp({'+_0 -_0': 1, '+_3 -_3': 1}, num_spin_orbitals=6)

import numpy as np
evolution_times = np.linspace(0.1, 5, 40)

from qiskit_cold_atom.applications import FermionicEvolutionProblem

spin_problem = FermionicEvolutionProblem(system, initial_state, evolution_times, spin_density)
particle_problem = FermionicEvolutionProblem(system, initial_state, evolution_times, particle_density)

from qiskit_cold_atom.providers.fermionic_tweezer_backend import FermionicTweezerSimulator
from qiskit_cold_atom.applications import TimeEvolutionSolver

# initializing the fermionic solver
fermionic_backend = FermionicTweezerSimulator(n_tweezers=3)
fermionic_solver = TimeEvolutionSolver(backend = fermionic_backend)

from qiskit import Aer

# initializing the qubit solver with the qasm simulator backend
qubit_backend = Aer.get_backend('qasm_simulator')

mapping = 'bravyi_kitaev'
# mapping = 'jordan_wigner'
# mapping = 'parity'

shallow_qubit_solver = TimeEvolutionSolver(backend = qubit_backend, map_type = mapping, trotter_steps = 1)
deep_qubit_solver = TimeEvolutionSolver(backend = qubit_backend, map_type = mapping, trotter_steps = 6)

spin_vals_fermions = fermionic_solver.solve(spin_problem)
particle_vals_fermions = fermionic_solver.solve(particle_problem)

spin_vals_qubits_shallow = shallow_qubit_solver.solve(spin_problem)
particle_vals_qubits_shallow = shallow_qubit_solver.solve(particle_problem)

spin_vals_qubits_deep = deep_qubit_solver.solve(spin_problem)
particle_vals_qubits_deep = deep_qubit_solver.solve(particle_problem)

import matplotlib.pyplot as plt

plt.rcParams['figure.figsize'] = [9, 6]
plt.rcParams.update({'font.size': 14})

plt.xlabel('evolution time')
plt.ylabel('Spin at site 1')
plt.ylim(-1.1, 1.1)

plt.title('evolution of spin density')
plt.plot(evolution_times, np.real_if_close(spin_vals_qubits_shallow), '--o', color='#d02670', alpha=0.5, label='qubits, 1 trotter step')
plt.plot(evolution_times, np.real_if_close(spin_vals_qubits_deep), '-o', color='#08bdba', alpha=0.8, label='qubits, 6 trotter steps')
plt.plot(evolution_times, np.real_if_close(spin_vals_fermions), '-o', color='#0f62fe', alpha=0.8, label='fermionic backend')
plt.legend()
plt.show()

plt.xlabel('evolution time')
plt.ylabel('Particle number at site 1')

plt.title('evolution of particle density')
plt.plot(evolution_times, particle_vals_qubits_shallow, '--o', color='#d02670', alpha=0.5, label='qubits, 1 trotter step')
plt.plot(evolution_times, particle_vals_qubits_deep, '-o', color='#08bdba', alpha=0.8, label='qubits, 6 trotter steps')
plt.plot(evolution_times, particle_vals_fermions, '-o', color='#0f62fe', alpha=0.8, label='fermionic backend')
plt.legend()
plt.show()

fermion_circuit = spin_problem.circuits(fermionic_backend.initialize_circuit(initial_state.occupations))[-1]
fermion_circuit.measure_all()

qubit_circuit = deep_qubit_solver.construct_qubit_circuits(spin_problem)[-1]
qubit_circuit.measure_all()

fermion_circuit.draw(output='mpl', style="clifford")

from qiskit import transpile
from qiskit.providers.fake_provider import FakeGuadalupe

qubit_device = FakeGuadalupe()

transpiled_circ = transpile(qubit_circuit, qubit_device, optimization_level=3)

transpiled_circ.draw(output='mpl', idle_wires=False, style="clifford")

import qiskit.tools.jupyter
%qiskit_version_table
%qiskit_copyright