"""High-level CPU expectation runner used by CLI scripts."""

from __future__ import annotations

import time
from dataclasses import dataclass

import numpy as np
from qibo.backends import construct_backend

from qibotn.benchmark_cases import exact_pauli_sum
from qibotn.observables import check_observable


@dataclass
class ExpectationConfig:
    ansatz: str = "tn"
    mpi: bool = False
    bond: int | None = 1024
    cut_ratio: float | None = 1e-12
    tensor_module: str = "torch"
    quimb_backend: str = "torch"
    dtype: str = "complex128"
    torch_threads: int = 8
    parallel_opts: dict | None = None


@dataclass
class ExpectationResult:
    value: float
    seconds: float
    rank: int = 0
    parallel_stats: list | None = None


def exact_for_observable(circuit, observable, nqubits):
    if isinstance(observable, dict) and "terms" in observable:
        terms = [
            (
                term["coefficient"],
                tuple((name, site) for name, site in term["operators"]),
            )
            for term in observable["terms"]
        ]
        return exact_pauli_sum(circuit, terms, nqubits)

    hamiltonian = check_observable(observable, nqubits)
    return float(hamiltonian.expectation_from_state(circuit().state(numpy=True)).real)


def run_cpu_expectation(circuit, observable, config):
    runcard = {
        "MPI_enabled": config.mpi,
        "MPS_enabled": config.ansatz.lower() == "mps",
        "NCCL_enabled": False,
        "expectation_enabled": observable,
        "max_bond_dimension": config.bond,
        "cut_ratio": config.cut_ratio,
        "tensor_module": config.tensor_module,
        "quimb_backend": config.quimb_backend,
        "dtype": config.dtype,
        "torch_threads": config.torch_threads,
        "parallel_opts": config.parallel_opts or {},
    }
    backend = construct_backend(
        backend="qibotn",
        platform="cpu",
        runcard=runcard,
    )

    start = time.perf_counter()
    value = backend.execute_circuit(circuit)[0]
    elapsed = time.perf_counter() - start

    rank = getattr(backend, "rank", 0)
    stats = getattr(backend, "parallel_stats", None)
    return ExpectationResult(
        float(np.real(value)),
        elapsed,
        rank=rank,
        parallel_stats=list(stats) if stats is not None else None,
    )
