from qiskit.circuit import Parameter
from qiskit import QuantumCircuit
from qiskit.primitives import Estimator
from qiskit.quantum_info import SparsePauliOp
from qiskit.algorithms import TimeEvolutionProblem
from qiskit.algorithms.time_evolvers import TrotterQRTE
from qiskit.synthesis import LieTrotter

import numpy as np

A = lambda t: 1 - t
B = lambda t: t
t_param = Parameter("t")

op1 = SparsePauliOp(["X"], np.array([2*A(t_param)]))
op2 = SparsePauliOp(["Z"], np.array([2*B(t_param)]))
op3 = SparsePauliOp(["Y"], np.array([1000]))
H_t = op1 + op2 + op3
print(H_t)

aux_op = [SparsePauliOp(["Z"])]

T = 2
reps = 10

init = QuantumCircuit(1)
# init.h(0)

evolution_problem = TimeEvolutionProblem(
    H_t,
    time=T,
    initial_state=init,
    aux_operators=aux_op,
    t_param=t_param
)
pf = LieTrotter(reps=1)
estimator = Estimator()

trotter_qrte = TrotterQRTE(product_formula=pf, num_timesteps=reps, estimator=estimator)
result = trotter_qrte.evolve(evolution_problem)

full_evo_circ = result.evolved_state

print(full_evo_circ.decompose().decompose())

aux_op_res = result.aux_ops_evaluated
obs_evo = result.observables

print(f"final result (previous only aux_op evaluation):  {aux_op_res}\n")
print("Newly added aux_op evaluations at each time step:")
for i, m in enumerate(obs_evo):
    print(f"step {i}:  <Z> = {m}")

import qiskit.tools.jupyter

%qiskit_version_table
%qiskit_copyright

