from qiskit_cold_atom.providers import ColdAtomProvider

provider = ColdAtomProvider()

for backend in provider.backends():
    print(backend)

import numpy as np
from qiskit import QuantumCircuit
backend = provider.get_backend("collective_spin_simulator")

circ_x = QuantumCircuit(1)
circ_x.rlx(np.pi/2, 0)
circ_x.measure_all()

circ_x.draw("mpl")

from qiskit.visualization import plot_histogram

job_x = backend.run(circ_x, shots = 1000, spin=20, seed=14)
plot_histogram(job_x.result().get_counts())

squeez_circ = QuantumCircuit(1, 1)
squeez_circ.rlx(-np.pi/2, 0)
squeez_circ.rlz2(0.3, 0)
squeez_circ.rlz(-np.pi/2, 0)
squeez_circ.rlx(-0.15, 0)
squeez_circ.measure(0, 0)

squeez_circ.draw(output='mpl')

job_squeez = backend.run(squeez_circ, shots = 1000, spin=20, seed=14)
plot_histogram(job_squeez.result().get_counts())

import matplotlib.pyplot as plt

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14,5))
fig.tight_layout(pad=5.0)
ax1.set_title("spin squeezing circuit")
ax2.set_title("single rotation circuit")

plot_histogram(job_x.result().get_counts(), ax=ax2)
plot_histogram(job_squeez.result().get_counts(), ax=ax1, color="#9f1853")

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

# One spin_up particle in the left site & one spin-down particle in the right site
initial_state = FermionicState([[1, 0, 0], [0, 0, 1]])

from qiskit_nature.operators.second_quantization import FermionicOp

spin_density = FermionicOp([('NIIIII', 1), ('IIINII', -1)])

evolution_times = np.linspace(0.1, 5, 40) # times to simulate

from qiskit_cold_atom.applications import FermionicEvolutionProblem

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

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

fermionic_backend = FermionicTweezerSimulator(n_tweezers=3)
fermionic_solver = TimeEvolutionSolver(backend = fermionic_backend)

from qiskit import Aer

qubit_backend = Aer.get_backend('qasm_simulator')

mapping = 'bravyi_kitaev'  # or 'jordan_wigner', '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)
spin_vals_qubits_shallow = shallow_qubit_solver.solve(spin_problem)
spin_vals_qubits_deep = deep_qubit_solver.solve(spin_problem)

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

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

qubit_circuit = deep_qubit_solver.construct_qubit_circuits(spin_problem)[-1]
qubit_circuit.measure_all()
qubit_circuit.decompose().decompose().draw("mpl", idle_wires=False)