"""Vidal/TEBD fast-path backend with qmatchatea fallback.

This backend targets MPS-friendly one-dimensional circuits: one-qubit gates and
adjacent two-qubit gates, measured with Pauli-sum expectation values. Unsupported
features fall back to the qmatchatea backend so the public behavior remains
usable while the fast path is expanded.
"""

from __future__ import annotations

import re
from dataclasses import dataclass

import numpy as np
from qibo.backends import NumpyBackend

from qibotn.backends.abstract import QibotnBackend
from qibotn.backends.qmatchatea import QMatchaTeaBackend
from qibotn.backends.vidal_mpi_segment import SegmentVidalMPIExecutor
from qibotn.backends.vidal_tebd import VidalTEBDExecutor, _gate_sites
from qibotn.observables import check_observable


def _symbolic_hamiltonian_to_pauli_terms(hamiltonian):
    terms = []
    factor_pattern = re.compile(r"([^\d]+)(\d+)")
    for term in hamiltonian.terms:
        ops = []
        for factor in term.factors:
            match = factor_pattern.match(str(factor))
            if match is None:
                raise ValueError(f"Unsupported observable factor {factor!r}.")
            name = match.group(1).upper()
            if name not in ("I", "X", "Y", "Z"):
                raise ValueError(f"Unsupported observable operator {name!r}.")
            if name != "I":
                ops.append((name, int(match.group(2))))
        terms.append((complex(term.coefficient), tuple(ops)))
    return terms


def _symbolic_hamiltonian_to_operator_terms(hamiltonian):
    terms = []
    factor_pattern = re.compile(r"([^\d]+)(\d+)")
    paulis = {
        "I": np.eye(2, dtype=np.complex128),
        "X": np.array([[0, 1], [1, 0]], dtype=np.complex128),
        "Y": np.array([[0, -1j], [1j, 0]], dtype=np.complex128),
        "Z": np.array([[1, 0], [0, -1]], dtype=np.complex128),
    }
    for term in hamiltonian.terms:
        ops_by_site = {}
        for factor in term.factors:
            site = getattr(factor, "target_qubit", None)
            matrix = getattr(factor, "matrix", None)
            if site is None or matrix is None:
                match = factor_pattern.match(str(factor))
                if match is None:
                    raise ValueError(f"Unsupported observable factor {factor!r}.")
                name = match.group(1).upper()
                if name not in paulis:
                    raise ValueError(f"Unsupported observable operator {name!r}.")
                site = int(match.group(2))
                matrix = paulis[name]
            matrix = np.asarray(matrix, dtype=np.complex128)
            site = int(site)
            if site in ops_by_site:
                ops_by_site[site] = ops_by_site[site] @ matrix
            else:
                ops_by_site[site] = matrix
        terms.append((complex(term.coefficient), tuple(ops_by_site.items())))
    return terms


def _dense_operator_to_product_terms(coeff, qubits, matrix):
    """Expand a dense k-local operator into product-matrix terms.

    The dense matrix basis is ordered by the provided ``qubits`` sequence.  For
    example, ``qubits=[2, 5]`` means matrix rows/columns are ordered as
    ``|q2 q5>``.
    """
    qubits = tuple(int(qubit) for qubit in qubits)
    if len(set(qubits)) != len(qubits):
        raise ValueError("Dense observable qubits must be unique.")
    matrix = np.asarray(matrix, dtype=np.complex128)
    dim = 2 ** len(qubits)
    if matrix.shape != (dim, dim):
        raise ValueError(
            "Dense observable matrix shape must be "
            f"({dim}, {dim}) for {len(qubits)} qubits."
        )

    units = [
        np.array([[1, 0], [0, 0]], dtype=np.complex128),
        np.array([[0, 1], [0, 0]], dtype=np.complex128),
        np.array([[0, 0], [1, 0]], dtype=np.complex128),
        np.array([[0, 0], [0, 1]], dtype=np.complex128),
    ]
    terms = []
    for row in range(dim):
        for col in range(dim):
            value = complex(coeff) * complex(matrix[row, col])
            if value == 0:
                continue
            ops = []
            for offset, site in enumerate(qubits):
                shift = len(qubits) - offset - 1
                out_bit = (row >> shift) & 1
                in_bit = (col >> shift) & 1
                ops.append((site, units[2 * out_bit + in_bit]))
            terms.append((value, tuple(ops)))
    return terms


def _dense_observable_to_operator_terms(observable):
    if not isinstance(observable, dict):
        return None

    if "matrix" in observable:
        terms = [observable]
    else:
        terms = observable.get("dense_terms")
        if terms is None:
            raw_terms = observable.get("terms")
            if not raw_terms or not any("matrix" in term for term in raw_terms):
                return None
            terms = raw_terms

    operator_terms = []
    for term in terms:
        if "matrix" not in term:
            raise ValueError("Dense observable terms must include a matrix.")
        qubits = term.get("qubits", term.get("sites"))
        if qubits is None:
            raise ValueError("Dense observable terms must include qubits or sites.")
        operator_terms.extend(
            _dense_operator_to_product_terms(
                term.get("coefficient", 1.0),
                qubits,
                term["matrix"],
            )
        )
    return operator_terms


def _operator_terms_to_mpo(terms, nqubits):
    """Build an exact direct-sum MPO for product-operator terms.

    This intentionally favors correctness and generality over compression: an
    ``m``-term sum becomes an MPO with bond dimension ``m``.  Local Hamiltonians
    can be compressed later without changing the public expectation path.
    """
    identity = np.eye(2, dtype=np.complex128)
    expanded_terms = []
    for coeff, ops in terms:
        local_ops = [identity for _ in range(nqubits)]
        for site, matrix in ops:
            site = int(site)
            if site < 0 or site >= nqubits:
                raise ValueError(f"Observable site {site} is outside the circuit.")
            matrix = np.asarray(matrix, dtype=np.complex128)
            if matrix.shape != (2, 2):
                raise ValueError("Only qubit local operators with shape (2, 2) are supported.")
            local_ops[site] = matrix
        expanded_terms.append((complex(coeff), local_ops))

    if not expanded_terms:
        raise ValueError("Cannot build an MPO from an empty observable.")

    bond_dim = len(expanded_terms)
    mpo = []
    for site in range(nqubits):
        left_dim = 1 if site == 0 else bond_dim
        right_dim = 1 if site == nqubits - 1 else bond_dim
        tensor = np.zeros((left_dim, 2, 2, right_dim), dtype=np.complex128)
        for term_index, (coeff, local_ops) in enumerate(expanded_terms):
            left = 0 if site == 0 else term_index
            right = 0 if site == nqubits - 1 else term_index
            op = coeff * local_ops[site] if site == 0 else local_ops[site]
            tensor[left, :, :, right] += op
        mpo.append(tensor)
    return mpo


def _observable_mpo_tensors(observable, nqubits=None):
    if isinstance(observable, dict):
        if "mpo_tensors" in observable:
            return observable["mpo_tensors"]
        if "mpo" in observable:
            return observable["mpo"]
        if nqubits is not None:
            terms = _dense_observable_to_operator_terms(observable)
            if terms is not None:
                return _operator_terms_to_mpo(terms, nqubits)
    return None


def _unsupported_reason(circuit):
    for gate in circuit.queue:
        name = getattr(gate, "name", gate.__class__.__name__)
        sites = _gate_sites(gate)
        if not sites:
            return f"gate {name} has no target qubits"
        if len(sites) > 2:
            return f"gate {name} acts on {len(sites)} qubits"
        if len(sites) == 2 and abs(sites[0] - sites[1]) != 1:
            return f"gate {name} is non-adjacent on qubits {sites}"
        if not hasattr(gate, "matrix"):
            return f"gate {name} does not expose a matrix"
    return None


def _can_route_non_adjacent(circuit):
    """True if the circuit's only unsupported feature is non-adjacent 2Q gates.

    SWAP routing can fix non-adjacent gates at compile time.  Multi-qubit
    gates and matrix-less gates are truly unsupported.
    """
    for gate in circuit.queue:
        sites = _gate_sites(gate)
        if not sites:
            return False
        if len(sites) > 2:
            return False
        if not hasattr(gate, "matrix"):
            return False
    return True


@dataclass
class _PreparedCircuit:
    nqubits: int
    queue: list


def _decompose_gate_for_mps(gate, nqubits, stack=()):
    sites = _gate_sites(gate)
    if len(sites) <= 2:
        return [gate]
    if gate in stack or not hasattr(gate, "decompose"):
        name = getattr(gate, "name", gate.__class__.__name__)
        raise ValueError(f"gate {name} acts on {len(sites)} qubits")

    free = [qubit for qubit in range(nqubits) if qubit not in sites]
    try:
        decomposed = gate.decompose(*free, use_toffolis=False, method="standard")
    except TypeError:
        decomposed = gate.decompose(*free)
    if not decomposed or decomposed == [gate]:
        name = getattr(gate, "name", gate.__class__.__name__)
        raise ValueError(f"gate {name} could not be decomposed for Vidal MPS")

    result = []
    for item in decomposed:
        result.extend(_decompose_gate_for_mps(item, nqubits, stack + (gate,)))
    return result


def _prepare_circuit_for_mps(circuit, decompose=True):
    if not decompose:
        return circuit
    queue = []
    for gate in circuit.queue:
        queue.extend(_decompose_gate_for_mps(gate, circuit.nqubits))
    return _PreparedCircuit(nqubits=circuit.nqubits, queue=queue)


@dataclass
class VidalBackend(QibotnBackend, NumpyBackend):
    """QiboTN backend using Vidal/TEBD when possible.

    The fast path supports:
    - one-qubit gates with ``gate.matrix()``;
    - adjacent two-qubit gates with ``gate.matrix()``;
    - Qibo ``SymbolicHamiltonian`` / qibotn dict Pauli-sum expectation values;
    - MPI chain segmentation through ``mpi_approach="CT"``.

    Unsupported operations are delegated to qmatchatea.
    """

    def __init__(self):
        super().__init__()
        self.name = "qibotn"
        self.platform = "vidal"
        self.precision = "double"
        self.rank = 0
        self.last_truncation_error = 0.0
        self.last_max_truncation_error = 0.0
        self.configure_tn_simulation()

    def configure_tn_simulation(
        self,
        ansatz: str = "MPS",
        max_bond_dimension: int | None = 10,
        cut_ratio: float | None = 1e-9,
        trunc_tracking_mode: str = "C",
        svd_control: str = "E!",
        ini_bond_dimension: int = 1,
        tensor_module: str = "torch",
        compile_circuit: bool = False,
        cache_gate_tensors: bool = True,
        track_memory: bool = False,
        mpi_approach: str = "SR",
        mpi_num_procs: int = 1,
        mpi_where_barriers: int = -1,
        mpi_isometrization: int = -1,
        mpi_term_batch_size: int | None = None,
        fallback: bool = True,
    ):
        self.ansatz = ansatz
        self.max_bond_dimension = max_bond_dimension
        self.cut_ratio = cut_ratio
        self.trunc_tracking_mode = trunc_tracking_mode
        self.svd_control = svd_control
        self.ini_bond_dimension = ini_bond_dimension
        self.tensor_module = tensor_module
        self.compile_circuit = compile_circuit
        self.cache_gate_tensors = cache_gate_tensors
        self.track_memory = track_memory
        self.mpi_approach = mpi_approach.upper()
        self.mpi_num_procs = mpi_num_procs
        self.mpi_where_barriers = mpi_where_barriers
        self.mpi_isometrization = mpi_isometrization
        self.mpi_term_batch_size = mpi_term_batch_size
        self.fallback = fallback
        self._fallback_backend = None

    def _setup_backend_specifics(self):
        return None

    def _qmatchatea_fallback(self):
        if self._fallback_backend is None:
            backend = QMatchaTeaBackend()
            backend.configure_tn_simulation(
                ansatz=self.ansatz,
                max_bond_dimension=self.max_bond_dimension,
                cut_ratio=self.cut_ratio,
                trunc_tracking_mode=self.trunc_tracking_mode,
                svd_control=self.svd_control,
                ini_bond_dimension=self.ini_bond_dimension,
                tensor_module=self.tensor_module,
                compile_circuit=self.compile_circuit,
                cache_gate_tensors=self.cache_gate_tensors,
                track_memory=self.track_memory,
                mpi_approach=self.mpi_approach,
                mpi_num_procs=self.mpi_num_procs,
                mpi_where_barriers=self.mpi_where_barriers,
                mpi_isometrization=self.mpi_isometrization,
            )
            self._fallback_backend = backend
        return self._fallback_backend

    def _fallback_or_raise(self, reason):
        if not self.fallback:
            raise NotImplementedError(reason)
        return self._qmatchatea_fallback()

    def _preprocess_circuit(self, circuit, compile_circuit):
        """Decompose unsupported multi-qubit gates for the local Vidal path."""
        return _prepare_circuit_for_mps(circuit, decompose=True)

    def _run_fast_executor(self, circuit, compile_circuit=True):
        if self.mpi_approach == "CT":
            from mpi4py import MPI

            self.rank = MPI.COMM_WORLD.Get_rank()
            executor = SegmentVidalMPIExecutor(
                nqubits=circuit.nqubits,
                max_bond=self.max_bond_dimension,
                cut_ratio=self.cut_ratio,
                tensor_module=self.tensor_module,
                comm=MPI.COMM_WORLD,
            )
        else:
            self.rank = 0
            executor = VidalTEBDExecutor(
                nqubits=circuit.nqubits,
                max_bond=self.max_bond_dimension,
                cut_ratio=self.cut_ratio,
                tensor_module=self.tensor_module,
            )
        executor.run_circuit(circuit, compile_circuit=compile_circuit)
        return executor

    def expectation(self, circuit, observable, preprocess=True, compile_circuit=None):
        if self.ansatz.upper() != "MPS":
            backend = self._fallback_or_raise("VidalBackend supports only MPS.")
            return backend.expectation(circuit, observable, preprocess, compile_circuit)

        original_circuit = circuit
        if compile_circuit is None:
            compile_circuit = self.compile_circuit

        if preprocess:
            try:
                circuit = self._preprocess_circuit(circuit, compile_circuit)
            except Exception as exc:
                backend = self._fallback_or_raise(
                    f"VidalBackend preprocessing failed: {exc}"
                )
                return backend.expectation(
                    original_circuit, observable, preprocess, compile_circuit
                )

        reason = _unsupported_reason(circuit)
        if reason is not None:
            # Non-adjacent gates can be routed at compile time
            if compile_circuit and _can_route_non_adjacent(circuit):
                pass  # proceed with Vidal + SWAP routing
            else:
                backend = self._fallback_or_raise(reason)
                return backend.expectation(
                    original_circuit, observable, preprocess, compile_circuit
                )

        executor = self._run_fast_executor(circuit, compile_circuit=compile_circuit)
        self.last_truncation_error = float(
            executor.global_truncation_error()
            if hasattr(executor, "global_truncation_error")
            else executor.truncation_error
        )
        self.last_max_truncation_error = float(
            executor.global_max_truncation_error()
            if hasattr(executor, "global_max_truncation_error")
            else executor.max_truncation_error
        )

        mpo_tensors = _observable_mpo_tensors(observable, circuit.nqubits)
        if mpo_tensors is not None:
            if self.mpi_approach == "CT":
                value = executor.expectation_mpo_root(mpo_tensors)
                from qtealeaves.tooling.mpisupport import MPI

                if MPI is not None and MPI.COMM_WORLD.Get_rank() != 0:
                    return np.nan
                return value
            return executor.expectation_mpo(mpo_tensors)

        hamiltonian = check_observable(observable, circuit.nqubits)
        try:
            terms = _symbolic_hamiltonian_to_operator_terms(hamiltonian)
        except ValueError as exc:
            backend = self._fallback_or_raise(str(exc))
            return backend.expectation(
                original_circuit, observable, preprocess, compile_circuit
            )

        mpo_tensors = _operator_terms_to_mpo(terms, circuit.nqubits)
        if self.mpi_approach == "CT":
            value = executor.expectation_mpo_root(mpo_tensors)
            from qtealeaves.tooling.mpisupport import MPI

            if MPI is not None and MPI.COMM_WORLD.Get_rank() != 0:
                return np.nan
            return value
        return executor.expectation_mpo(mpo_tensors)

    def execute_circuit(
        self,
        circuit,
        initial_state=None,
        nshots=None,
        prob_type=None,
        return_array=False,
        **prob_kwargs,
    ):
        backend = self._fallback_or_raise(
            "VidalBackend.execute_circuit is delegated to qmatchatea."
        )
        return backend.execute_circuit(
            circuit,
            initial_state=initial_state,
            nshots=nshots,
            prob_type=prob_type,
            return_array=return_array,
            **prob_kwargs,
        )
