"""CPU tensor-network backend with cutensornet-style runcard support."""

from __future__ import annotations

import os
import pickle
import time
from pathlib import Path

import numpy as np
from qibo import hamiltonians
from qibo.backends import NumpyBackend
from qibo.config import raise_error

from qibotn.backends.abstract import QibotnBackend
from qibotn.backends.vidal import (
    _observable_mpo_tensors,
    _operator_terms_to_mpo,
    _symbolic_hamiltonian_to_operator_terms,
    _unsupported_reason,
)
from qibotn.backends.vidal_mpi_segment import SegmentVidalMPIExecutor
from qibotn.backends.vidal_tebd import VidalTEBDExecutor
from qibotn.observables import check_observable
from qibotn.result import TensorNetworkResult


def _as_bool_or_dict(value, name):
    if isinstance(value, (bool, dict)):
        return value
    raise TypeError(f"{name} has an unexpected type")


def _bind_numa_node(rank):
    """Bind the calling process (or thread) to the NUMA node for *rank*.

    The MPI rank is converted to a local (per-node) rank through the
    environment variables commonly set by Open MPI, MVAPICH, and Slurm.
    The process CPU affinity and NUMA memory policy are set accordingly.

    Returns the NUMA domain that was selected, or ``None`` if the binding
    could not be determined.
    """
    current_affinity = os.sched_getaffinity(0)
    online_cpus = set(range(os.cpu_count() or 1))
    if current_affinity and current_affinity != online_cpus:
        # MPI launchers such as Intel MPI often pin local ranks correctly
        # before Python starts.  Do not narrow that placement further.
        return None

    local_rank = rank
    for name in (
        "OMPI_COMM_WORLD_LOCAL_RANK",
        "MV2_COMM_WORLD_LOCAL_RANK",
        "MPI_LOCALRANKID",
        "I_MPI_LOCAL_RANK",
        "SLURM_LOCALID",
    ):
        try:
            local_rank = int(os.environ[name])
            break
        except (KeyError, ValueError):
            pass

    domains = _available_numa_domains()
    if not domains:
        return None

    local_size = _local_world_size()
    assigned_domains = domains[local_rank::local_size]
    if not assigned_domains:
        assigned_domains = [domains[local_rank % len(domains)]]

    domain = assigned_domains[0]
    cpus = set()
    for selected in assigned_domains:
        cpulist = f"/sys/devices/system/node/node{selected}/cpulist"
        try:
            cpus.update(_parse_cpu_list(open(cpulist, encoding="utf-8").read().strip()))
        except (FileNotFoundError, OSError):
            pass
    try:
        if cpus:
            os.sched_setaffinity(0, cpus)
    except OSError:
        pass

    try:
        import ctypes

        libnuma = ctypes.CDLL("libnuma.so.1")
        if libnuma.numa_available() >= 0:
            libnuma.numa_run_on_node(ctypes.c_int(domain))
            libnuma.numa_set_preferred(ctypes.c_int(domain))
    except Exception:
        pass

    return domain


def _available_numa_domains():
    nodes = []
    base = Path("/sys/devices/system/node")
    try:
        for path in base.glob("node[0-9]*"):
            try:
                nodes.append(int(path.name[4:]))
            except ValueError:
                pass
    except OSError:
        return []
    return sorted(nodes)


def _local_world_size():
    for name in (
        "OMPI_COMM_WORLD_LOCAL_SIZE",
        "MV2_COMM_WORLD_LOCAL_SIZE",
        "MPI_LOCALNRANKS",
        "I_MPI_LOCAL_SIZE",
        "SLURM_NTASKS_PER_NODE",
    ):
        value = os.environ.get(name)
        if not value:
            continue
        try:
            return max(1, int(str(value).split("(", 1)[0]))
        except ValueError:
            pass
    return 1


def _parse_cpu_list(text):
    cpus = set()
    for item in text.split(","):
        item = item.strip()
        if not item:
            continue
        if "-" in item:
            start, stop = item.split("-", 1)
            cpus.update(range(int(start), int(stop) + 1))
        else:
            cpus.add(int(item))
    return cpus


class CpuTensorNet(QibotnBackend, NumpyBackend):
    """CPU replacement for the cutensornet runcard execution surface.

    The backend preserves the high-level runcard knobs used by the GPU backend:
    ``MPI_enabled``, ``MPS_enabled`` and ``expectation_enabled``.  Generic TN
    work is delegated to quimb on CPU; MPS expectation uses the Vidal fast path
    when the circuit is nearest-neighbor and falls back to quimb otherwise.
    """

    def __init__(self, runcard=None):
        super().__init__()
        self.name = "qibotn"
        self.platform = "cpu"
        self.precision = "double"
        self.configure_tn_simulation(runcard)

    def configure_tn_simulation(self, runcard=None):
        runcard = {} if runcard is None else runcard
        self.rank = 0
        self.MPI_enabled = bool(runcard.get("MPI_enabled", False))
        self.NCCL_enabled = bool(runcard.get("NCCL_enabled", False))
        if self.NCCL_enabled:
            raise_error(NotImplementedError, "NCCL is only available for GPU backends.")

        expectation = runcard.get("expectation_enabled", False)
        if expectation is True:
            self.expectation_enabled = True
            self.observable = None
        elif expectation is False:
            self.expectation_enabled = False
            self.observable = None
        elif isinstance(expectation, (dict, hamiltonians.SymbolicHamiltonian)):
            self.expectation_enabled = True
            self.observable = expectation
        else:
            raise TypeError("expectation_enabled has an unexpected type")

        mps = _as_bool_or_dict(runcard.get("MPS_enabled", False), "MPS_enabled")
        self.MPS_enabled = bool(mps)
        self.mps_options = mps if isinstance(mps, dict) else {}

        self.max_bond_dimension = runcard.get(
            "max_bond_dimension",
            self.mps_options.get("max_bond_dimension", 512),
        )
        self.cut_ratio = runcard.get(
            "cut_ratio",
            self.mps_options.get(
                "cut_ratio",
                self.mps_options.get("svd_method", {}).get("abs_cutoff", 1e-12),
            ),
        )
        self.tensor_module = runcard.get("tensor_module", "torch")
        self.dtype = runcard.get("dtype", "complex128")
        self.compile_circuit = bool(runcard.get("compile_circuit", False))
        self.preprocess = bool(runcard.get("preprocess", False))
        self.mpi_term_batch_size = runcard.get(
            "mpi_term_batch_size",
            runcard.get("term_batch_size", None),
        )
        self.torch_threads = runcard.get("torch_threads", None)
        self.quimb_backend = runcard.get("quimb_backend", "torch")
        self.contraction_optimizer = runcard.get("contraction_optimizer", "auto-hq")
        self.parallel_opts = runcard.get("parallel_opts", {})
        self.parallel_stats = []

    def execute_circuit(
        self,
        circuit,
        initial_state=None,
        nshots=None,
        prob_type=None,
        return_array=False,
        **prob_kwargs,
    ):
        if initial_state is not None:
            raise_error(NotImplementedError, "QiboTN CPU backend does not support initial state.")

        if self.torch_threads is not None and self.tensor_module == "torch":
            import torch

            torch.set_num_threads(self.torch_threads)

        if self.expectation_enabled:
            value = self.expectation(circuit, self.observable)
            if self.MPI_enabled and self.rank > 0:
                return np.asarray([0], dtype=np.int64)
            dtype = np.complex128 if np.iscomplexobj(value) else np.float64
            return np.asarray([value], dtype=dtype)

        backend = self._quimb_backend()
        backend.configure_tn_simulation(
            ansatz="mps" if self.MPS_enabled else None,
            max_bond_dimension=self.max_bond_dimension if self.MPS_enabled else None,
            svd_cutoff=self.cut_ratio,
        )
        return backend.execute_circuit(
            circuit=circuit,
            nshots=nshots,
            return_array=return_array,
        )

    def expectation(self, circuit, observable=None, preprocess=None, compile_circuit=None):
        mpo_tensors = _observable_mpo_tensors(observable, circuit.nqubits)
        if mpo_tensors is None:
            observable = check_observable(observable, circuit.nqubits)
        use_preprocess = self.preprocess if preprocess is None else preprocess
        if mpo_tensors is not None and not self.MPS_enabled:
            raise_error(
                NotImplementedError,
                "MPO expectation is currently supported only by the Vidal MPS path.",
            )

        if self.MPS_enabled:
            reason = _unsupported_reason(circuit)
            if reason is None or self.compile_circuit or use_preprocess:
                return self._vidal_expectation(
                    circuit,
                    observable,
                    preprocess=use_preprocess,
                    compile_circuit=compile_circuit,
                )

        backend = self._quimb_backend()
        backend.configure_tn_simulation(
            ansatz="mps" if self.MPS_enabled else None,
            max_bond_dimension=self.max_bond_dimension if self.MPS_enabled else None,
            svd_cutoff=self.cut_ratio,
        )
        if self.MPI_enabled:
            return self._quimb_expectation_mpi(backend, circuit, observable)
        return self._quimb_expectation_processpool(backend, circuit, observable)

    def _vidal_expectation(
        self, circuit, observable, preprocess=False, compile_circuit=None
    ):
        if compile_circuit is None:
            compile_circuit = self.compile_circuit
        if preprocess:
            if self.MPI_enabled:
                from mpi4py import MPI

                self.rank = MPI.COMM_WORLD.Get_rank()

            from qibotn.backends.vidal import VidalBackend

            backend = VidalBackend()
            backend.configure_tn_simulation(
                max_bond_dimension=self.max_bond_dimension,
                cut_ratio=self.cut_ratio,
                tensor_module=self.tensor_module,
                compile_circuit=compile_circuit,
                mpi_approach="CT" if self.MPI_enabled else "SR",
                mpi_term_batch_size=self.mpi_term_batch_size,
                fallback=False,
            )
            value = backend.expectation(
                circuit,
                observable,
                preprocess=True,
                compile_circuit=compile_circuit,
            )
            self.rank = getattr(backend, "rank", self.rank)
            self.last_truncation_error = getattr(
                backend, "last_truncation_error", np.nan
            )
            self.last_max_truncation_error = getattr(
                backend, "last_max_truncation_error", np.nan
            )
            return value

        mpo_tensors = _observable_mpo_tensors(observable, circuit.nqubits)
        if self.MPI_enabled:
            from mpi4py import MPI

            comm = MPI.COMM_WORLD
            self.rank = comm.Get_rank()
            executor = SegmentVidalMPIExecutor(
                nqubits=circuit.nqubits,
                max_bond=self.max_bond_dimension,
                cut_ratio=self.cut_ratio,
                tensor_module=self.tensor_module,
                comm=comm,
            )
            executor.run_circuit(circuit)
            self.last_truncation_error = float(executor.global_truncation_error())
            self.last_max_truncation_error = float(
                executor.global_max_truncation_error()
            )
            if mpo_tensors is not None:
                value = executor.expectation_mpo_root(mpo_tensors)
            else:
                terms = _symbolic_hamiltonian_to_operator_terms(observable)
                value = executor.expectation_mpo_root(
                    _operator_terms_to_mpo(terms, circuit.nqubits)
                )
            return np.nan if self.rank != 0 else value

        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)
        self.last_truncation_error = float(executor.truncation_error)
        self.last_max_truncation_error = float(executor.max_truncation_error)
        if mpo_tensors is not None:
            return executor.expectation_mpo(mpo_tensors)
        terms = _symbolic_hamiltonian_to_operator_terms(observable)
        return executor.expectation_mpo(_operator_terms_to_mpo(terms, circuit.nqubits))

    def _quimb_backend(self):
        import qibotn.backends.quimb as qmb

        return qmb.BACKENDS[self.quimb_backend](
            quimb_backend=self.quimb_backend,
            contraction_optimizer=self.contraction_optimizer,
        )

    def _bind_rank_to_numa_domain(self, rank):
        self.numa_domain = _bind_numa_node(rank)

    def _default_search_workers(self, nranks=1):
        if self.torch_threads:
            return max(1, int(self.torch_threads))
        return max(1, (os.cpu_count() or 1) // max(1, nranks))

    def _quimb_expectation_processpool(self, backend, circuit, observable):
        return self._quimb_expectation_search(
            backend,
            circuit,
            observable,
            method="processpool",
            comm=None,
        )

    def _quimb_expectation_mpi(self, backend, circuit, observable):
        from mpi4py import MPI

        comm = MPI.COMM_WORLD
        self.rank = comm.Get_rank()
        self._bind_rank_to_numa_domain(self.rank)

        return self._quimb_expectation_search(
            backend,
            circuit,
            observable,
            method="mpi",
            comm=comm,
        )

    def _quimb_expectation_search(self, backend, circuit, observable, method, comm=None):
        rank = comm.Get_rank() if comm is not None else 0
        size = comm.Get_size() if comm is not None else 1
        self.rank = rank

        from qibotn.observables import extract_gates_and_qubits
        from qibotn.parallel import (
            contraction_tree_costs,
            parallel_contract,
            parallel_path_search,
        )
        from qibotn.backends.quimb import (
            PAULI_DENSE_MAX_QUBITS,
            _pauli_term_to_dense_operator,
            pauli_product_expectation_tn,
        )

        opts = dict(self.parallel_opts)
        user_slicing_opts = opts.get("slicing_opts")
        search_workers = opts.get("search_workers", self._default_search_workers(size))
        search_repeats = opts.get("max_repeats", 128)
        search_time = opts.get("max_time", 60)
        search_backend = opts.get("search_backend")
        dask_address = opts.get("dask_address")
        dask_close_workers = bool(opts.get("dask_close_workers", False))
        print_stats = bool(opts.get("print_stats", False))
        debug_trials = bool(opts.get("debug_trials", False))
        search_only = bool(opts.get("search_only", False))
        save_tree_path = opts.get("save_tree_path")
        load_tree_path = opts.get("load_tree_path")
        loaded_trees = None
        saved_trees = []
        saved_costs = []

        if load_tree_path:
            with Path(load_tree_path).open("rb") as f:
                payload = pickle.load(f)
            loaded_trees = payload["trees"] if isinstance(payload, dict) else payload
            if not isinstance(loaded_trees, (list, tuple)):
                loaded_trees = [loaded_trees]

        qc = backend._qibo_circuit_to_quimb(
            circuit,
            quimb_circuit_type=backend.circuit_ansatz,
            gate_opts={
                "max_bond": self.max_bond_dimension,
                "cutoff": self.cut_ratio,
            },
        )

        total_value = 0.0 + 0.0j
        terms = extract_gates_and_qubits(observable)
        for term_index, (coeff, factors) in enumerate(terms):
            if not factors:
                if self.rank == 0:
                    total_value += coeff
                continue

            if len(factors) > PAULI_DENSE_MAX_QUBITS:
                tn = pauli_product_expectation_tn(
                    qc,
                    factors,
                    simplify_sequence="ADCRS",
                    simplify_atol=1e-12,
                )
            else:
                op, where = _pauli_term_to_dense_operator(factors)
                tn = qc.local_expectation(
                    op,
                    where,
                    rehearse="tn",
                    simplify_sequence="ADCRS",
                    simplify_atol=1e-12,
                )
            slicing_opts = self._mpi_slicing_opts(
                user_slicing_opts,
            )

            if loaded_trees is not None:
                if term_index >= len(loaded_trees):
                    raise ValueError(
                        f"Loaded tree file has {len(loaded_trees)} tree(s), "
                        f"but term {term_index} was requested."
                    )
                tree = loaded_trees[term_index]
                search_seconds = 0.0
                if self.rank == 0 and print_stats:
                    print(
                        f"tn_tree_loaded term={term_index} path={load_tree_path}",
                        flush=True,
                    )
            else:
                search_start = time.perf_counter()
                tree = parallel_path_search(
                    tn,
                    tn.outer_inds(),
                    method="dask" if method != "mpi" and search_backend == "dask" else method,
                    total_repeats=search_repeats,
                    max_time=search_time,
                    n_workers=search_workers,
                    slicing_opts=slicing_opts,
                    trial_timeout=opts.get("trial_timeout"),
                    search_backend=search_backend,
                    dask_address=dask_address,
                    debug_trials=debug_trials,
                    dask_close_workers=dask_close_workers,
                )
                search_seconds = time.perf_counter() - search_start
            if tree is None:
                raise RuntimeError("Failed to find a contraction tree for CPU TN MPI.")
            if self.parallel_opts.get("contract_implementation") == "cpp":
                from qibotn.torch_contractor import prepare_torch_cpp_contractor

                prepare_torch_cpp_contractor(tree)

            path_cost = contraction_tree_costs(tree)
            search_stats = getattr(tree, "qibotn_search_stats", {})
            if save_tree_path and loaded_trees is None:
                saved_trees.append(tree)
                saved_costs.append(path_cost)
            if self.rank == 0 and print_stats:
                print(
                    "tn_search_done "
                    f"term={term_index} "
                    f"search_seconds={search_seconds:.3f} "
                    f"completed_trials={search_stats.get('completed_trials', 'na')} "
                    f"finite_trials={search_stats.get('finite_trials', 'na')} "
                    f"failed_trials={search_stats.get('failed_trials', 'na')} "
                    f"requested_trials={search_stats.get('requested_trials', search_repeats)} "
                    f"best_score={search_stats.get('best_score', float('nan')):.6g} "
                    f"slices={path_cost['nslices']} "
                    f"log10_flops={path_cost['log10_flops']:.3f} "
                    f"log10_write={path_cost['log10_write']:.3f} "
                    f"log2_size={path_cost['log2_size']:.3f} "
                    f"log10_combo={path_cost['log10_combo']:.3f} "
                    f"peak_memory_gib={path_cost['peak_memory_gib']:.6g}",
                    flush=True,
                )

            if search_only:
                self.parallel_stats.append(
                    {
                        "term_index": term_index,
                        "term_factors": tuple(factors),
                        "path_cost": path_cost,
                        "search_stats": search_stats,
                        "tree_slices": int(getattr(tree, "multiplicity", 1)),
                        "slice_assignment": "search_only",
                        "rank_slices": [],
                        "search_seconds": search_seconds,
                        "contract_seconds": 0.0,
                        "search_workers": search_workers,
                        "search_repeats": search_repeats,
                        "search_time": search_time,
                        "search_backend": search_backend or method,
                        "dask_address": dask_address,
                        "numa_domain": getattr(self, "numa_domain", None),
                    }
                )
                continue

            if comm is None and int(getattr(tree, "multiplicity", 1)) <= 1:
                if self.rank == 0:
                    contract_start = time.perf_counter()
                    value = self._contract_term_unsliced(tn, tree, backend)
                    contract_seconds = time.perf_counter() - contract_start
                    if print_stats:
                        print(
                            "tn_contract_done "
                            f"term={term_index} "
                            f"contract_seconds={contract_seconds:.3f}",
                            flush=True,
                        )
                    self.parallel_stats.append(
                        {
                            "term_index": term_index,
                            "term_factors": tuple(factors),
                            "path_cost": path_cost,
                            "search_stats": search_stats,
                            "tree_slices": 1,
                            "slice_assignment": "root",
                            "rank_slices": [1] + [0] * (size - 1),
                            "search_seconds": search_seconds,
                            "contract_seconds": contract_seconds,
                            "search_workers": search_workers,
                            "search_repeats": search_repeats,
                            "search_time": search_time,
                            "search_backend": search_backend or method,
                            "dask_address": dask_address,
                            "numa_domain": getattr(self, "numa_domain", None),
                        }
                    )
                    total_value += coeff * complex(value)
                continue

            if comm is None:
                contract_start = time.perf_counter()
                value = self._contract_term_unsliced(tn, tree, backend)
                contract_seconds = time.perf_counter() - contract_start
                if print_stats:
                    print(
                        "tn_contract_done "
                        f"term={term_index} "
                        f"contract_seconds={contract_seconds:.3f}",
                        flush=True,
                    )
                self.parallel_stats.append(
                    {
                        "term_index": term_index,
                        "term_factors": tuple(factors),
                        "path_cost": path_cost,
                        "search_stats": search_stats,
                        "tree_slices": int(getattr(tree, "multiplicity", 1)),
                        "slice_assignment": "local",
                        "rank_slices": [int(getattr(tree, "multiplicity", 1))],
                        "search_seconds": search_seconds,
                        "contract_seconds": contract_seconds,
                        "search_workers": search_workers,
                        "search_repeats": search_repeats,
                        "search_time": search_time,
                        "search_backend": search_backend or method,
                        "dask_address": dask_address,
                        "numa_domain": getattr(self, "numa_domain", None),
                    }
                )
                total_value += coeff * complex(np.asarray(value).reshape(-1)[0])
                continue

            contract_start = time.perf_counter()
            arrays = self._term_arrays(tn, backend)
            value, stats = parallel_contract(
                tree,
                arrays,
                method="mpi",
                comm=comm,
                return_stats=True,
                implementation=self.parallel_opts.get("contract_implementation"),
            )
            contract_seconds = time.perf_counter() - contract_start
            gathered_stats = comm.gather(stats, root=0)
            if rank == 0:
                if print_stats:
                    print(
                        "tn_contract_done "
                        f"term={term_index} "
                        f"contract_seconds={contract_seconds:.3f}",
                        flush=True,
                    )
                self.parallel_stats.append(
                    {
                        "term_index": term_index,
                        "term_factors": tuple(factors),
                        "path_cost": path_cost,
                        "search_stats": search_stats,
                        "tree_slices": stats.nslices,
                        "slice_assignment": stats.assignment,
                        "rank_slices": [
                            item.local_slices for item in gathered_stats
                        ],
                        "search_seconds": search_seconds,
                        "contract_seconds": contract_seconds,
                        "search_workers": search_workers,
                        "search_repeats": search_repeats,
                        "search_time": search_time,
                        "search_backend": search_backend or method,
                        "dask_address": dask_address,
                        "numa_domain": getattr(self, "numa_domain", None),
                    }
                )
                total_value += coeff * complex(np.asarray(value).reshape(-1)[0])

        if self.rank == 0 and save_tree_path and loaded_trees is None:
            path = Path(save_tree_path)
            path.parent.mkdir(parents=True, exist_ok=True)
            with path.open("wb") as f:
                pickle.dump(
                    {
                        "trees": saved_trees,
                        "costs": saved_costs,
                        "nterms": len(saved_trees),
                    },
                    f,
                    protocol=pickle.HIGHEST_PROTOCOL,
                )
            if print_stats:
                print(
                    f"tn_tree_saved path={save_tree_path} nterms={len(saved_trees)}",
                    flush=True,
                )

        if search_only:
            return np.nan

        return np.nan if rank != 0 else float(np.real(total_value))

    def _contract_term_unsliced(self, tn, tree, backend):
        contract_implementation = self.parallel_opts.get("contract_implementation")
        if contract_implementation == "cpp":
            if backend.backend != "torch":
                raise ValueError("contract_implementation='cpp' requires torch backend.")
            from qibotn.backends.quimb import _torch_cpu_array, _torch_dtype
            from qibotn.torch_contractor import contract_tree_cpp

            arrays = [
                _torch_cpu_array(array, dtype=_torch_dtype(self.dtype))
                for array in tn.arrays
            ]
            nslices = int(getattr(tree, "multiplicity", 1))
            if nslices > 1:
                total = None
                for slice_id in range(nslices):
                    value = contract_tree_cpp(tree, tree.slice_arrays(arrays, slice_id))
                    total = value if total is None else total + value
                return total
            return contract_tree_cpp(tree, arrays)

        if backend.backend == "torch":
            from qibotn.backends.quimb import _torch_cpu_array, _torch_dtype

            for tensor in tn.tensors:
                tensor._data = _torch_cpu_array(
                    tensor._data,
                    dtype=_torch_dtype(self.dtype),
                )
            return tn.contract(
                all,
                output_inds=(),
                optimize=tree,
                backend="torch",
                implementation=contract_implementation,
            )

        return tn.contract(
            all,
            output_inds=(),
            optimize=tree,
            backend=backend.backend,
            implementation=contract_implementation,
        )

    def _mpi_slicing_opts(self, user_slicing_opts):
        return None if user_slicing_opts is None else dict(user_slicing_opts)

    def _term_arrays(self, tn, backend):
        if backend.backend == "torch":
            from qibotn.backends.quimb import _torch_cpu_array, _torch_dtype

            return [
                _torch_cpu_array(array, dtype=_torch_dtype(self.dtype))
                for array in tn.arrays
            ]
        from qibotn.backends.quimb import _numpy_dtype

        return [backend.engine.asarray(array, dtype=_numpy_dtype(self.dtype)) for array in tn.arrays]
