# SPDX-FileCopyrightText: 2025 Qoro Quantum Ltd <divi@qoroquantum.de>
#
# SPDX-License-Identifier: Apache-2.0

from functools import lru_cache

import numpy as np
import numpy.typing as npt
import pennylane as qml


def _get_structural_key(obs: qml.operation.Operation) -> tuple[str, ...]:
    """Generates a hashable, wire-independent key from an observable's structure.

    This function is used to create a canonical representation of an observable
    based on its constituent Pauli operators, ignoring the wires they act on.
    This key is ideal for caching computed eigenvalues, as observables with the
    same structure (e.g., PauliX(0) @ PauliZ(1) and PauliX(2) @ PauliZ(3))
    share the same eigenvalues. It maps PauliX and PauliY to PauliZ because
    they are all isospectral (have eigenvalues [1, -1]).

    Args:
        obs (qml.operation.Operation): A PennyLane observable (e.g., qml.PauliZ(0), qml.PauliX(0) @ qml.PauliY(1)).

    Returns:
        tuple[str, ...]: A tuple of strings representing the structure of the observable,
            e.g., ('PauliZ',) or ('PauliZ', 'PauliZ').
    """

    # Pennylane returns the same eigenvalues for PauliX and PauliY
    # since it handles diagonalizing gates internally anyway
    name_map = {
        "PauliY": "PauliZ",
        "PauliX": "PauliZ",
        "PauliZ": "PauliZ",
        "Identity": "Identity",
    }

    if isinstance(obs, qml.ops.Prod):
        # Recursively build a tuple of operator names
        return tuple(name_map[o.name] for o in obs.operands)

    # For single operators, return a single-element tuple
    return (name_map[obs.name],)


@lru_cache(maxsize=512)
def _get_eigvals_from_key(key: tuple[str, ...]) -> npt.NDArray[np.int8]:
    """Computes and caches eigenvalues based on a structural key.

    This function takes a key generated by `_get_structural_key` and computes
    the eigenvalues of the corresponding tensor product of operators. The results
    are memoized using @lru_cache to avoid redundant calculations.

    Args:
        key (tuple[str, ...]): A tuple of strings representing the observable's structure.

    Returns:
        np.ndarray: A NumPy array containing the eigenvalues of the observable.
    """

    # Define a mapping from name to the base eigenvalue array
    eigvals_map = {
        "PauliZ": np.array([1, -1], dtype=np.int8),
        "Identity": np.array([1, 1], dtype=np.int8),
    }

    # Start with the eigenvalues of the first operator in the key
    final_eigvals = eigvals_map[key[0]]

    # Iteratively compute the kronecker product for the rest
    for op_name in key[1:]:
        final_eigvals = np.kron(final_eigvals, eigvals_map[op_name])

    return final_eigvals


def _batched_expectation(
    shots_dicts: list[dict[str, int]],
    observables: list[qml.operation.Operation],
    wire_order: tuple[int, ...],
) -> npt.NDArray[np.float64]:
    """Efficiently calculates expectation values for multiple observables across multiple shot histograms.

    This function is optimized to compute expectation values in a fully vectorized
    manner, minimizing Python loops. It operates in four main steps:
    1. Aggregates all unique bitstrings measured across all histograms.
    2. Builds a "reduced" eigenvalue matrix corresponding only to the unique states.
    3. Builds a "reduced" probability matrix from the shot counts for each histogram.
    4. Computes all expectation values with a single matrix multiplication.

    Args:
        shots_dicts (list[dict[str, int]]): A list of shot dictionaries (histograms),
            where each dictionary maps a measured bitstring to its count.
        observables (list[qml.operation.Operation]): A list of PennyLane observables
            for which to calculate expectation values.
        wire_order (tuple[int, ...]): A tuple defining the order of wires, which maps
            the bitstring to the qubits. Note: This is typically the reverse of the
            qubit indices (e.g., (2, 1, 0) for a 3-qubit system).

    Returns:
        npt.NDArray[np.float64]: A 2D NumPy array of shape (n_observables, n_shots) where
            result[i, j] is the expectation value of observables[i] for the
            histogram in shots_dicts[j].
    """

    n_histograms = len(shots_dicts)
    n_total_wires = len(wire_order)
    n_observables = len(observables)

    # --- 1. Aggregate all unique measured states across all shots ---
    all_measured_bitstrings = set()
    for sd in shots_dicts:
        all_measured_bitstrings.update(sd.keys())

    unique_bitstrings = sorted(list(all_measured_bitstrings))
    n_unique_states = len(unique_bitstrings)

    bitstring_to_idx_map = {bs: i for i, bs in enumerate(unique_bitstrings)}

    # --- 2. Build REDUCED Eigenvalue Matrix: (n_observables, n_unique_states) ---
    # For systems with <=64 qubits, use fast integer conversion and bitwise operations.
    # For larger systems, use character array to avoid integer overflow.
    if n_total_wires <= 64:
        # Fast path: convert to uint64 and use vectorized bitwise operations
        unique_states_int = np.array(
            [int(bs, 2) for bs in unique_bitstrings], dtype=np.uint64
        )
        use_integer_representation = True
    else:
        # Safe path: convert to character array for large qubit counts
        bitstring_chars = np.array([list(bs) for bs in unique_bitstrings], dtype="U1")
        use_integer_representation = False

    reduced_eigvals_matrix = np.zeros((n_observables, n_unique_states))
    wire_map = {w: i for i, w in enumerate(wire_order)}

    powers_cache = {}

    for obs_idx, observable in enumerate(observables):
        obs_wires = observable.wires
        n_obs_wires = len(obs_wires)

        if n_obs_wires in powers_cache:
            powers = powers_cache[n_obs_wires]
        else:
            powers = 2 ** np.arange(n_obs_wires - 1, -1, -1, dtype=np.intp)
            powers_cache[n_obs_wires] = powers

        obs_wire_indices = np.array([wire_map[w] for w in obs_wires], dtype=np.uint32)
        eigvals = _get_eigvals_from_key(_get_structural_key(observable))

        # Vectorized mapping, but on the *reduced* set of states
        shifts = n_total_wires - 1 - obs_wire_indices

        if use_integer_representation:
            # Fast path: vectorized bitwise operations on integers
            bits = ((unique_states_int[:, np.newaxis] >> shifts) & 1).astype(np.intp)
        else:
            # Safe path: extract bits from character array (vectorized)
            bits = bitstring_chars[:, shifts].astype(np.intp)

        obs_state_indices = np.dot(bits, powers)

        reduced_eigvals_matrix[obs_idx, :] = eigvals[obs_state_indices]

    # --- 3. Build REDUCED Probability Matrix: (n_shots, n_unique_states) ---
    reduced_prob_matrix = np.zeros((n_histograms, n_unique_states), dtype=np.float32)
    for i, shots_dict in enumerate(shots_dicts):
        total = sum(shots_dict.values())

        for bitstring, count in shots_dict.items():
            col_idx = bitstring_to_idx_map[bitstring]
            reduced_prob_matrix[i, col_idx] = count / total

    # --- 4. Compute Final Expectation Values ---
    # (n_shots, n_unique_states) @ (n_unique_states, n_observables)
    result = reduced_prob_matrix @ reduced_eigvals_matrix.T

    # Transpose to (n_observables, n_shots) as expected by the calling code
    return result.T
