"""Segmented MPI Vidal/TEBD executor.

Each rank owns a contiguous interval of sites. Gates fully inside an interval
are applied locally. Only two-site gates crossing a rank boundary communicate
the neighboring edge tensor and the resulting boundary update.
"""

from __future__ import annotations

import time
from dataclasses import dataclass

import numpy as np
from mpi4py import MPI

from qibotn.backends.vidal_tebd import (
    _asarray,
    _backend_module,
    _build_theta_svd_matrix,
    _disjoint_batches,
    _fuse_one_site_blocks,
    _gate_sites,
    _is_two_qubit_batch,
    _make_two_site_update,
    _ones,
    _real_if_close,
    _route_non_adjacent_gates,
    _svd,
    _tensor_update_from_numpy,
    _tensor_update_to_numpy,
    _to_float,
    _to_numpy,
    _transpose,
    VidalTEBDExecutor,
)

_EDGE_TAG = 1701
_UPDATE_TAG = 1702
_EXPECT_ENV_TAG = 1703
_EXPECT_RESULT_TAG = 1704


def _partition_sites(nsites, nranks):
    base = nsites // nranks
    rem = nsites % nranks
    starts = [0]
    for rank in range(nranks):
        starts.append(starts[-1] + base + int(rank < rem))
    return starts


@dataclass
class SegmentVidalMPIExecutor:
    nqubits: int
    max_bond: int | None
    comm: object
    cut_ratio: float | None = 1e-12
    tensor_module: str = "torch"

    def __post_init__(self):
        self.rank = self.comm.Get_rank()
        self.size = self.comm.Get_size()
        self.starts = _partition_sites(self.nqubits, self.size)
        self.start = self.starts[self.rank]
        self.end = self.starts[self.rank + 1]
        if self.start == self.end:
            raise ValueError("SegmentVidalMPIExecutor requires at least one site per rank.")

        from qibotn.backends.cpu import _bind_numa_node

        self.numa_domain = _bind_numa_node(self.rank)

        self.xp = _backend_module(self.tensor_module)
        if self.xp is np:
            self.dtype = np.complex128
            self.device = None
        else:
            self.dtype = self.xp.complex128
            self.device = self.xp.device("cpu")

        self.gammas = {}
        for site in range(self.start, self.end):
            self.gammas[site] = _asarray(
                self.xp, [[[1.0 + 0.0j], [0.0 + 0.0j]]], self.dtype
            )

        self.lambdas = {
            bond: _ones(self.xp, 1, self.dtype, self.device)
            for bond in range(self.start, self.end + 1)
        }
        self._accumulated_truncation_error = 0.0
        self._max_truncation_error = 0.0

    @property
    def truncation_error(self):
        return self._accumulated_truncation_error

    def global_truncation_error(self):
        return self.comm.allreduce(self._accumulated_truncation_error, op=MPI.SUM)

    @property
    def max_truncation_error(self):
        return self._max_truncation_error

    def global_max_truncation_error(self):
        return self.comm.allreduce(self._max_truncation_error, op=MPI.MAX)

    def owns_site(self, site):
        return self.start <= site < self.end

    def owner_of(self, site):
        return int(np.searchsorted(self.starts, site, side="right") - 1)

    def run_circuit(self, circuit, compile_circuit=True):
        timings = {
            "local_compute": 0.0,
            "edge_exchange": 0.0,
            "boundary_compute": 0.0,
            "boundary_update": 0.0,
            "one_site": 0.0,
            "gather": 0.0,
        }

        gates = circuit.queue
        if compile_circuit:
            gates = _route_non_adjacent_gates(gates, circuit.nqubits)
        gates = _fuse_one_site_blocks(gates)
        for batch in _disjoint_batches(gates):
            if _is_two_qubit_batch(batch):
                self._apply_two_site_batch(batch, timings)
            else:
                tic = time.perf_counter()
                for gate in batch:
                    sites = _gate_sites(gate)
                    if len(sites) == 1 and self.owns_site(sites[0]):
                        op = _asarray(self.xp, gate.matrix(), self.dtype)
                        self.apply_one_site(op, sites[0])
                    elif len(sites) == 2:
                        self._apply_two_site_batch([gate], timings)
                    elif len(sites) > 2:
                        raise NotImplementedError("Only one- and two-qubit gates are supported.")
                timings["one_site"] += time.perf_counter() - tic

        return timings

    def apply_one_site(self, op, pos):
        self.gammas[pos] = self.xp.einsum("st,atb->asb", op, self.gammas[pos])

    def _apply_two_site_batch(self, batch, timings):
        local_gates = []
        boundary_specs = []
        recv_left_update = False
        for gate in batch:
            sites = _gate_sites(gate)
            if abs(sites[0] - sites[1]) != 1:
                raise NotImplementedError("Segment Vidal supports adjacent two-qubit gates only.")
            left, right = sorted(sites)
            left_owner = self.owner_of(left)
            right_owner = self.owner_of(right)
            if left_owner == self.rank and right_owner == self.rank:
                local_gates.append(gate)
            elif left_owner == self.rank:
                boundary_specs.append((gate, left, right))
            elif right_owner == self.rank:
                recv_left_update = True

        tic = time.perf_counter()
        edge_send_req = None
        if recv_left_update:
            edge_send_req = self.comm.isend(
                self._edge_payload(), dest=self.rank - 1, tag=_EDGE_TAG
            )
        right_edge = (
            self.comm.recv(source=self.rank + 1, tag=_EDGE_TAG)
            if boundary_specs
            else None
        )
        timings["edge_exchange"] += time.perf_counter() - tic

        boundary_update = None
        tic = time.perf_counter()
        for gate, left, right in boundary_specs:
            boundary_update = self._compute_boundary_update(
                gate, left, right, right_edge
            )
        timings["boundary_compute"] += time.perf_counter() - tic

        tic = time.perf_counter()
        update_send_req = None
        if boundary_update is not None:
            update_send_req = self.comm.isend(
                boundary_update, dest=self.rank + 1, tag=_UPDATE_TAG
            )
        timings["boundary_update"] += time.perf_counter() - tic

        tic = time.perf_counter()
        local_items = [
            self._compute_owned_two_site_update(gate)
            for gate in local_gates
        ]
        timings["local_compute"] += time.perf_counter() - tic

        tic = time.perf_counter()
        left_boundary_update = (
            self.comm.recv(source=self.rank - 1, tag=_UPDATE_TAG)
            if recv_left_update
            else None
        )
        if update_send_req is not None:
            update_send_req.wait()
        if edge_send_req is not None:
            edge_send_req.wait()
        timings["boundary_update"] += time.perf_counter() - tic

        for update in local_items:
            self._install_update(update)
        if boundary_update is not None:
            self._install_update(boundary_update)
        if left_boundary_update is not None:
            self._install_update(left_boundary_update)

    def _edge_payload(self):
        return {
            "start": self.start,
            "end": self.end,
            "gamma_start": _to_numpy(self.gammas[self.start]),
            "lambda_after_start": _to_numpy(self.lambdas[self.start + 1]),
        }

    def _compute_owned_two_site_update(self, gate):
        sites = _gate_sites(gate)
        op = _asarray(self.xp, gate.matrix(), self.dtype)
        left, right = sites
        if left > right:
            left, right = right, left
            op = _transpose(self.xp, op.reshape(2, 2, 2, 2), (1, 0, 3, 2)).reshape(4, 4)
        item = self._build_item(
            left,
            op,
            self.lambdas[left],
            self.lambdas[left + 1],
            self.lambdas[left + 2],
            self.gammas[left],
            self.gammas[right],
        )
        split = _svd(self.xp, item["matrix"])
        return _make_two_site_update(
            item, *split, self.max_bond, self.cut_ratio, self.xp
        )

    def _compute_boundary_update(self, gate, left, right, remote):
        op = _asarray(self.xp, gate.matrix(), self.dtype)
        sites = _gate_sites(gate)
        if sites[0] > sites[1]:
            op = _transpose(self.xp, op.reshape(2, 2, 2, 2), (1, 0, 3, 2)).reshape(4, 4)

        gamma_right = _asarray(self.xp, remote["gamma_start"], self.dtype)
        lam_right = _asarray(
            self.xp,
            remote["lambda_after_start"],
            self.xp.float64 if self.xp is not np else np.float64,
        )
        item = self._build_item(
            left,
            op,
            self.lambdas[left],
            self.lambdas[left + 1],
            lam_right,
            self.gammas[left],
            gamma_right,
        )
        split = _svd(self.xp, item["matrix"])
        return _tensor_update_to_numpy(
            _make_two_site_update(
                item, *split, self.max_bond, self.cut_ratio, self.xp
            )
        )

    def _build_item(self, site, op, lam_left, lam_mid, lam_right, gamma_left, gamma_right):
        result = _build_theta_svd_matrix(
            op, self.xp, lam_left, lam_mid, lam_right, gamma_left, gamma_right
        )
        result["site"] = site
        result["lam_left"] = lam_left
        result["lam_right"] = lam_right
        return result

    def _install_update(self, update):
        if isinstance(update["left"], np.ndarray):
            update = _tensor_update_from_numpy(self.xp, update, self.dtype)
        truncation_error = update.get("truncation_error", 0.0)
        self._accumulated_truncation_error += truncation_error
        self._max_truncation_error = max(
            self._max_truncation_error,
            truncation_error,
        )
        site = update["site"]
        if self.owns_site(site):
            self.gammas[site] = update["left"]
        if self.owns_site(site + 1):
            self.gammas[site + 1] = update["right"]
        if self.start <= site + 1 <= self.end:
            self.lambdas[site + 1] = update["lambda"]

    def gather_full_state(self):
        payload = {
            "start": self.start,
            "end": self.end,
            "gammas": {site: _to_numpy(tensor) for site, tensor in self.gammas.items()},
            "lambdas": {bond: _to_numpy(tensor) for bond, tensor in self.lambdas.items()},
        }
        return self.comm.gather(payload, root=0)

    def expectation_pauli_sum_root(self, terms, term_batch_size=None):
        paulis = {
            "I": self._eye(2),
            "X": _asarray(self.xp, [[0, 1], [1, 0]], self.dtype),
            "Y": _asarray(self.xp, [[0, -1j], [1j, 0]], self.dtype),
            "Z": _asarray(self.xp, [[1, 0], [0, -1]], self.dtype),
        }
        operator_terms = [
            (
                coeff,
                tuple((site, paulis[name.upper()]) for name, site in ops),
            )
            for coeff, ops in terms
        ]
        return self.expectation_operator_sum_root(
            operator_terms,
            term_batch_size=term_batch_size,
        )

    def expectation_operator_sum_root(self, terms, term_batch_size=None):
        if term_batch_size is None:
            term_batch_size = max(1, len(terms))
        norm = self._distributed_product_expectation({})
        total = 0.0 + 0.0j
        for start in range(0, len(terms), int(term_batch_size)):
            batch = terms[start : start + int(term_batch_size)]
            values = self._distributed_operator_batch_expectation(batch, norm)
            if self.rank == 0:
                for (coeff, _), term_value in zip(batch, values):
                    total += complex(coeff) * complex(term_value)
        return None if self.rank != 0 else _real_if_close(total / norm)

    def _eye(self, size):
        if self.xp is np:
            return np.eye(size, dtype=self.dtype)
        return self.xp.eye(size, dtype=self.dtype, device=self.device)

    def _distributed_product_expectation(self, operators):
        if self.rank == 0:
            env = self._segment_product_environment(operators)
            if self.size == 1:
                return env.reshape(-1)[0]
            self.comm.send(_to_numpy(env), dest=1, tag=_EXPECT_ENV_TAG)
            return self.comm.recv(source=self.size - 1, tag=_EXPECT_RESULT_TAG)

        incoming = self.comm.recv(source=self.rank - 1, tag=_EXPECT_ENV_TAG)
        env = self._segment_product_environment(operators, incoming)
        if self.rank == self.size - 1:
            self.comm.send(_to_numpy(env).reshape(-1)[0], dest=0, tag=_EXPECT_RESULT_TAG)
        else:
            self.comm.send(_to_numpy(env), dest=self.rank + 1, tag=_EXPECT_ENV_TAG)
        return None

    def _segment_product_environment(self, operators, incoming=None):
        if incoming is None:
            env = _asarray(
                self.xp,
                np.eye(len(self.lambdas[self.start]), dtype=np.complex128),
                self.dtype,
            )
        else:
            env = _asarray(self.xp, incoming, self.dtype)

        identity = self._eye(2)
        for site in range(self.start, self.end):
            tensor = self.gammas[site] * self.lambdas[site + 1].reshape(1, 1, -1)
            op = operators.get(site, identity)
            env = self.xp.einsum(
                "xy,xsb,st,ytd->bd", env, self._conj(tensor), op, tensor
            )
        return env

    def _distributed_operator_batch_expectation(self, terms, norm):
        if not terms:
            return []
        if all(not ops for _, ops in terms):
            return [norm] * len(terms) if self.rank == 0 else None

        batch_ops = [
            {int(site): _asarray(self.xp, matrix, self.dtype) for site, matrix in ops}
            for _, ops in terms
        ]
        if self.rank == 0:
            env = self._segment_operator_batch_environment(batch_ops)
            if self.size == 1:
                return list(env.reshape(len(terms), -1)[:, 0])
            self.comm.send(_to_numpy(env), dest=1, tag=_EXPECT_ENV_TAG)
            return self.comm.recv(source=self.size - 1, tag=_EXPECT_RESULT_TAG)

        incoming = self.comm.recv(source=self.rank - 1, tag=_EXPECT_ENV_TAG)
        env = self._segment_operator_batch_environment(batch_ops, incoming)
        if self.rank == self.size - 1:
            values = list(_to_numpy(env).reshape(len(terms), -1)[:, 0])
            self.comm.send(values, dest=0, tag=_EXPECT_RESULT_TAG)
        else:
            self.comm.send(_to_numpy(env), dest=self.rank + 1, tag=_EXPECT_ENV_TAG)
        return None

    def _segment_operator_batch_environment(self, batch_ops, incoming=None):
        batch_size = len(batch_ops)
        if incoming is None:
            dim = len(self.lambdas[self.start])
            env = _asarray(
                self.xp,
                np.tile(np.eye(dim, dtype=np.complex128), (batch_size, 1, 1)),
                self.dtype,
            )
        else:
            env = _asarray(self.xp, incoming, self.dtype)

        identity = self._eye(2)
        for site in range(self.start, self.end):
            tensor = self.gammas[site] * self.lambdas[site + 1].reshape(1, 1, -1)
            ops = self.xp.stack(
                [operators.get(site, identity) for operators in batch_ops],
                axis=0,
            )
            env = self.xp.einsum(
                "nxy,xsb,nst,ytd->nbd",
                env,
                self._conj(tensor),
                ops,
                tensor,
            )
        return env

    def _conj(self, tensor):
        return np.conjugate(tensor) if self.xp is np else tensor.conj()

    def expectation_mpo_root(self, mpo_tensors):
        if len(mpo_tensors) != self.nqubits:
            raise ValueError(
                f"Expected {self.nqubits} MPO tensors, got {len(mpo_tensors)}."
            )
        norm = self._distributed_product_expectation({})
        if self.rank == 0:
            env = self._segment_mpo_environment(mpo_tensors)
            if self.size == 1:
                return _real_if_close(env.reshape(-1)[0] / norm)
            self.comm.send(_to_numpy(env), dest=1, tag=_EXPECT_ENV_TAG)
            value = self.comm.recv(source=self.size - 1, tag=_EXPECT_RESULT_TAG)
            return _real_if_close(value / norm)

        incoming = self.comm.recv(source=self.rank - 1, tag=_EXPECT_ENV_TAG)
        env = self._segment_mpo_environment(mpo_tensors, incoming)
        if self.rank == self.size - 1:
            self.comm.send(
                _to_numpy(env).reshape(-1)[0],
                dest=0,
                tag=_EXPECT_RESULT_TAG,
            )
        else:
            self.comm.send(_to_numpy(env), dest=self.rank + 1, tag=_EXPECT_ENV_TAG)
        return None

    def _segment_mpo_environment(self, mpo_tensors, incoming=None):
        if incoming is None:
            left_dim = len(self.lambdas[self.start])
            env = _asarray(
                self.xp,
                np.zeros((left_dim, 1, left_dim), dtype=np.complex128),
                self.dtype,
            )
            env[:, 0, :] = self._eye(left_dim)
        else:
            env = _asarray(self.xp, incoming, self.dtype)

        for site in range(self.start, self.end):
            mpo = _asarray(self.xp, mpo_tensors[site], self.dtype)
            if mpo.ndim != 4 or mpo.shape[1:3] != (2, 2):
                raise ValueError(
                    "Each MPO tensor must have shape "
                    "(left_bond, 2, 2, right_bond)."
                )
            tensor = self.gammas[site] * self.lambdas[site + 1].reshape(1, 1, -1)
            env = self.xp.einsum(
                "xlc,xub,lutr,ctd->brd",
                env,
                self._conj(tensor),
                mpo,
                tensor,
            )
        return env

    def expectation_ring_xz_root(self):
        terms = [
            (0.5, (("X", site), ("Z", (site + 1) % self.nqubits)))
            for site in range(self.nqubits)
        ]
        return self.expectation_pauli_sum_root(terms)


def run_segment_vidal_mpi_ring_xz(
    circuit,
    max_bond,
    comm,
    cut_ratio=1e-12,
    tensor_module="torch",
):
    executor = SegmentVidalMPIExecutor(
        nqubits=circuit.nqubits,
        max_bond=max_bond,
        cut_ratio=cut_ratio,
        tensor_module=tensor_module,
        comm=comm,
    )
    timings = executor.run_circuit(circuit)
    tic = time.perf_counter()
    value = executor.expectation_ring_xz_root()
    timings["gather"] = time.perf_counter() - tic
    return value, timings
