"""Expectation value benchmark with timing breakdown."""
import time
import numpy as np
import torch
import qibo
from qibo import Circuit, gates, hamiltonians
from qibo.symbols import X, Z
import quimb as qu
from qibotn.parallel import parallel_path_search

# Config
NQUBITS = 25
NLAYERS = 10
WORKERS = 96

# Build circuit
np.random.seed(42)
circuit = Circuit(NQUBITS)
for _ in range(NLAYERS):
    for q in range(NQUBITS):
        circuit.add(gates.RY(q, theta=np.random.uniform(0, 2*np.pi)))
        circuit.add(gates.RZ(q, theta=np.random.uniform(0, 2*np.pi)))
    for q in range(NQUBITS):
        circuit.add(gates.CNOT(q % NQUBITS, (q + 1) % NQUBITS))

# Setup backend
qibo.set_backend("qibotn", platform="quimb")
backend = qibo.get_backend()
backend.configure_tn_simulation(ansatz="tn")
torch.set_num_threads(WORKERS)

# Prepare TN
t0 = time.time()
qc = backend._qibo_circuit_to_quimb(circuit, backend.circuit_ansatz)
op = qu.pauli('x') & qu.pauli('z')
tn = qc.local_expectation(op, (0, 1), rehearse='tn')
t_prepare = time.time() - t0

# Search path
t0 = time.time()
tree = parallel_path_search(tn, tn.outer_inds(), 'processpool', 1024, 300, WORKERS)
t_search = time.time() - t0

# Contract
t0 = time.time()
for tensor in tn.tensors:
    tensor._data = torch.from_numpy(np.asarray(tensor._data)).to(torch.complex128)
val = tn.contract(all, output_inds=(), optimize=tree)
t_contract = time.time() - t0

# Results
print(f"Prepare:  {t_prepare:.4f}s")
print(f"Search:   {t_search:.4f}s")
print(f"Contract: {t_contract:.4f}s")
print(f"Total:    {t_prepare+t_search+t_contract:.4f}s")
print(f"Expectation: {0.5 * complex(val).real:.10f}")