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

import logging
from abc import abstractmethod
from functools import lru_cache, partial
from itertools import chain, groupby
from queue import Queue

import numpy as np
import pennylane as qml
from scipy.optimize import OptimizeResult

from divi.backends import CircuitRunner
from divi.circuits import Circuit, MetaCircuit
from divi.circuits.qem import _NoMitigation
from divi.qprog.exceptions import _CancelledError
from divi.qprog.optimizers import ScipyMethod, ScipyOptimizer
from divi.qprog.quantum_program import QuantumProgram
from divi.reporting import LoggingProgressReporter, QueueProgressReporter
from divi.utils import hamiltonian_to_pauli_string, reverse_dict_endianness

logger = logging.getLogger(__name__)


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, ...]) -> np.ndarray:
    """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, observables, wire_order):
    """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:
        np.ndarray: 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) ---
    unique_states_int = np.array(
        [int(bs, 2) for bs in unique_bitstrings], dtype=np.uint64
    )
    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
        bits = ((unique_states_int[:, np.newaxis] >> shifts) & 1).astype(np.intp)
        # powers = 2 ** np.arange(n_obs_wires - 1, -1, -1)

        # obs_state_indices = (bits * powers).sum(axis=1).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


def _compute_parameter_shift_mask(n_params):
    """
    Generate a binary matrix mask for the parameter shift rule.
    This mask is used to determine the shifts to apply to each parameter
    when computing gradients via the parameter shift rule in quantum algorithms.

    Args:
        n_params (int): The number of parameters in the quantum circuit.

    Returns:
        np.ndarray: A (2 * n_params, n_params) matrix where each row encodes
            the shift to apply to each parameter for a single evaluation.
            The values are multiples of 0.5 * pi, with alternating signs.
    """
    mask_arr = np.arange(0, 2 * n_params, 2)
    mask_arr[0] = 1

    binary_matrix = ((mask_arr[:, np.newaxis] & (1 << np.arange(n_params))) > 0).astype(
        np.float64
    )

    binary_matrix = binary_matrix.repeat(2, axis=0)
    binary_matrix[1::2] *= -1
    binary_matrix *= 0.5 * np.pi

    return binary_matrix


class VariationalQuantumAlgorithm(QuantumProgram):
    """Base class for variational quantum algorithms.

    This class provides the foundation for implementing variational quantum
    algorithms in Divi. It handles circuit execution, parameter optimization,
    and result management for algorithms that optimize parameterized quantum
    circuits to minimize cost functions.

    Variational algorithms work by:
    1. Generating parameterized quantum circuits
    2. Executing circuits on quantum hardware/simulators
    3. Computing expectation values of cost Hamiltonians
    4. Using classical optimizers to update parameters
    5. Iterating until convergence

    Attributes:
        _losses_history (list[dict]): History of loss values during optimization.
        _final_params (np.ndarray): Final optimized parameters.
        _best_params (np.ndarray): Parameters that achieved the best loss.
        _best_loss (float): Best loss achieved during optimization.
        _circuits (list[Circuit]): Generated quantum circuits.
        _total_circuit_count (int): Total number of circuits executed.
        _total_run_time (float): Total execution time in seconds.
        _curr_params (np.ndarray): Current parameter values.
        _seed (int | None): Random seed for parameter initialization.
        _rng (np.random.Generator): Random number generator.
        _grad_mode (bool): Whether currently computing gradients.
        _grouping_strategy (str): Strategy for grouping quantum operations.
        _qem_protocol (QEMProtocol): Quantum error mitigation protocol.
        _cancellation_event (Event | None): Event for graceful termination.
        _meta_circuit_factory (callable): Factory for creating MetaCircuit instances.
    """

    def __init__(
        self,
        backend: CircuitRunner,
        seed: int | None = None,
        progress_queue: Queue | None = None,
        **kwargs,
    ):
        """Initialize the VariationalQuantumAlgorithm.

        This constructor is specifically designed for hybrid quantum-classical
        variational algorithms. The instance variables `n_layers` and `n_params`
        must be set by subclasses, where:
        - `n_layers` is the number of layers in the quantum circuit.
        - `n_params` is the number of parameters per layer.

        For exotic variational algorithms where these variables may not be applicable,
        the `_initialize_params` method should be overridden to set the parameters.

        Args:
            backend (CircuitRunner): Quantum circuit execution backend.
            seed (int | None): Random seed for parameter initialization. Defaults to None.
            progress_queue (Queue | None): Queue for progress reporting. Defaults to None.

        Keyword Args:
            grouping_strategy (str): Strategy for grouping operations in Pennylane transforms.
                Options: "default", "wires", "qwc". Defaults to "qwc".
            qem_protocol (QEMProtocol | None): Quantum error mitigation protocol to apply. Defaults to None.
        """

        self._losses_history = []

        self._best_params = None
        self._final_params = None
        self._best_loss = float("inf")
        self._best_probs = {}

        self._curr_params = None

        self._seed = seed
        self._rng = np.random.default_rng(self._seed)

        # Lets child classes adapt their optimization
        # step for grad calculation routine
        self._grad_mode = False
        self._is_compute_probabilites = False

        super().__init__(
            backend=backend, seed=seed, progress_queue=progress_queue, **kwargs
        )

        self.job_id = kwargs.get("job_id", None)
        if progress_queue and self.job_id:
            self.reporter = QueueProgressReporter(self.job_id, progress_queue)
        else:
            self.reporter = LoggingProgressReporter()

        # Needed for Pennylane's transforms
        self._grouping_strategy = kwargs.pop("grouping_strategy", "qwc")

        self._qem_protocol = kwargs.pop("qem_protocol", None) or _NoMitigation()

        self._cancellation_event = None

        self._meta_circuit_factory = partial(
            MetaCircuit,
            # No grouping strategy for expectation value measurements
            grouping_strategy=(
                "_backend_expval"
                if self.backend and self.backend.supports_expval
                else self._grouping_strategy
            ),
            qem_protocol=self._qem_protocol,
        )

    @property
    @abstractmethod
    def cost_hamiltonian(self) -> qml.operation.Operator:
        """The cost Hamiltonian for the variational problem."""
        pass

    @property
    def total_circuit_count(self) -> int:
        """Get the total number of circuits executed.

        Returns:
            int: Cumulative count of circuits submitted for execution.
        """
        return self._total_circuit_count

    @property
    def total_run_time(self) -> float:
        """Get the total runtime across all circuit executions.

        Returns:
            float: Cumulative execution time in seconds.
        """
        return self._total_run_time

    @property
    def meta_circuits(self):
        """Get the meta-circuit templates used by this program.

        Returns:
            dict[str, MetaCircuit]: Dictionary mapping circuit names to their
                MetaCircuit templates.
        """
        return self._meta_circuits

    @property
    def n_params(self):
        """Get the total number of parameters in the quantum circuit.

        Returns:
            int: Total number of trainable parameters (n_layers * n_params_per_layer).
        """
        return self._n_params

    @property
    def losses_history(self) -> list[dict]:
        """Get a copy of the optimization loss history.

        Each entry is a dictionary mapping parameter indices to loss values.

        Returns:
            list[dict]: Copy of the loss history. Modifications to this list
                will not affect the internal state.
        """
        return self._losses_history.copy()

    @property
    def final_params(self) -> np.ndarray:
        """Get a copy of the final optimized parameters.

        Returns:
            np.ndarray: Copy of the final parameters. Modifications to this array
                will not affect the internal state.
        """
        return self._final_params.copy()

    @property
    def best_params(self) -> np.ndarray:
        """Get a copy of the parameters that achieved the best (lowest) loss.

        Returns:
            np.ndarray: Copy of the best parameters. Modifications to this array
                will not affect the internal state.
        """
        return self._best_params.copy()

    @property
    def best_loss(self) -> float:
        """Get the best loss achieved so far.

        Returns:
            float: The best loss achieved so far.
        """
        return self._best_loss

    @property
    def best_probs(self):
        """Get a copy of the probability distribution for the best parameters.

        Returns:
            dict: A copy of the best probability distribution.
        """
        return self._best_probs.copy()

    def _convert_counts_to_probs(self, results):
        """Convert raw counts to probability distributions."""
        return {
            outer_k: {
                inner_k: inner_v / self.backend.shots
                for inner_k, inner_v in outer_v.items()
            }
            for outer_k, outer_v in results.items()
        }

    def _process_probability_results(self, results):
        """Convert raw counts to probabilities and fix endianness."""
        probs = self._convert_counts_to_probs(results)
        return reverse_dict_endianness(probs)

    @property
    def initial_params(self) -> np.ndarray:
        """Get the current initial parameters.

        Returns:
            np.ndarray: Current initial parameters. If not yet initialized,
                they will be generated automatically.
        """
        if self._curr_params is None:
            self._initialize_params()
        return self._curr_params.copy()

    @initial_params.setter
    def initial_params(self, value: np.ndarray | None):
        """
        Set initial parameters.

        Args:
            value (np.ndarray | None): Initial parameters with shape
                (n_param_sets, n_layers * n_params), or None to reset
                to uninitialized state.

        Raises:
            ValueError: If parameters have incorrect shape.
        """
        if value is not None:
            self._validate_initial_params(value)
            self._curr_params = value.copy()
        else:
            # Reset to uninitialized state
            self._curr_params = None

    @abstractmethod
    def _create_meta_circuits_dict(self) -> dict[str, MetaCircuit]:
        pass

    @abstractmethod
    def _generate_circuits(self, **kwargs) -> list[Circuit]:
        """Generate quantum circuits for execution.

        This method should generate and return a list of Circuit objects based on
        the current algorithm state and parameters. The circuits will be executed
        by the backend.

        Args:
            **kwargs: Additional keyword arguments for circuit generation.

        Returns:
            list[Circuit]: List of Circuit objects to be executed.
        """
        pass

    def get_expected_param_shape(self) -> tuple[int, int]:
        """
        Get the expected shape for initial parameters.

        Returns:
            tuple[int, int]: Shape (n_param_sets, n_layers * n_params) that
                initial parameters should have for this quantum program.
        """
        return (self.optimizer.n_param_sets, self.n_layers * self.n_params)

    def _validate_initial_params(self, params: np.ndarray):
        """
        Validate user-provided initial parameters.

        Args:
            params (np.ndarray): Parameters to validate.

        Raises:
            ValueError: If parameters have incorrect shape.
        """
        expected_shape = self.get_expected_param_shape()

        if params.shape != expected_shape:
            raise ValueError(
                f"Initial parameters must have shape {expected_shape}, "
                f"got {params.shape}"
            )

    def _initialize_params(self):
        """
        Initialize the circuit parameters randomly.

        Generates random parameters with values uniformly distributed between
        0 and 2π. The number of parameter sets depends on the optimizer being used.
        """
        total_params = self.n_layers * self.n_params
        self._curr_params = self._rng.uniform(
            0, 2 * np.pi, (self.optimizer.n_param_sets, total_params)
        )

    def _run_optimization_circuits(self, data_file, **kwargs):
        self._curr_circuits = self._generate_circuits(**kwargs)

        if self.backend.supports_expval:
            kwargs["ham_ops"] = hamiltonian_to_pauli_string(
                self.cost_hamiltonian, self.n_qubits
            )

        losses = self._dispatch_circuits_and_process_results(
            data_file=data_file, **kwargs
        )

        return losses

    def _post_process_results(
        self, results: dict[str, dict[str, int]], **kwargs
    ) -> dict[int, float]:
        """
        Post-process the results of the quantum problem.

        Args:
            results (dict[str, dict[str, int]]): The shot histograms of the quantum execution step.
                The keys should be strings of format {param_id}_*_{measurement_group_id}.
                i.e. an underscore-separated bunch of metadata, starting always with
                the index of some parameter and ending with the index of some measurement group.
                Any extra piece of metadata that might be relevant to the specific
                application can be kept in the middle.

        Returns:
            dict[int, float]: The energies for each parameter set grouping, where the dict keys
                correspond to the parameter indices.
        """
        if not (self._cancellation_event and self._cancellation_event.is_set()):
            self.reporter.info(
                message="Post-processing output", iteration=self.current_iteration
            )

        losses = {}
        measurement_groups = self._meta_circuits["cost_circuit"].measurement_groups

        # Flatten measurement groups for expectation value backends
        if self.backend.supports_expval:
            flattened_measurement_groups = tuple(
                chain.from_iterable(measurement_groups)
            )
        else:
            flattened_measurement_groups = measurement_groups

        # Define key functions for both levels of grouping
        get_param_id = lambda item: int(item[0].split("_")[0])
        get_qem_id = lambda item: int(item[0].split("_")[1].split(":")[1])

        # Group the pre-sorted results by parameter ID.
        for p, param_group_iterator in groupby(results.items(), key=get_param_id):
            param_group_iterator = list(param_group_iterator)

            shots_by_qem_idx = zip(
                *{
                    gid: [value for _, value in group]
                    for gid, group in groupby(param_group_iterator, key=get_qem_id)
                }.values()
            )

            marginal_results = []
            for shots_dicts, curr_measurement_group in zip(
                shots_by_qem_idx, flattened_measurement_groups
            ):
                if self.backend.supports_expval:
                    ham_ops = kwargs.get("ham_ops")
                    if ham_ops is None:
                        # Internal consistency check: ham_ops should be set by _run_optimization_circuits
                        # when backend supports expectation values
                        raise ValueError(
                            "Hamiltonian operators (ham_ops) are required when using a backend "
                            "that supports expectation values, but were not provided."
                        )

                    expectation_matrix = np.array(
                        [
                            [shot_dict[op_name] for op_name in ham_ops.split(";")]
                            for shot_dict in shots_dicts
                        ]
                    ).T
                else:
                    wire_order = tuple(reversed(self.cost_hamiltonian.wires))

                    expectation_matrix = _batched_expectation(
                        shots_dicts, curr_measurement_group, wire_order
                    )

                # expectation_matrix[i, j] = expectation value for observable i, histogram j
                curr_marginal_results = []
                for intermediate_exp_values in expectation_matrix:
                    mitigated_exp_value = self._qem_protocol.postprocess_results(
                        intermediate_exp_values
                    )
                    curr_marginal_results.append(mitigated_exp_value)

                marginal_results.append(
                    curr_marginal_results
                    if len(curr_marginal_results) > 1
                    else curr_marginal_results[0]
                )

            if self.backend.supports_expval:
                marginal_results = marginal_results[0]

            pl_loss = (
                self._meta_circuits["cost_circuit"]
                .postprocessing_fn(marginal_results)[0]
                .item()
            )

            losses[p] = pl_loss + self.loss_constant

        return losses

    def _perform_final_computation(self, **kwargs):
        """
        Perform final computations after optimization completes.

        This is an optional hook method that subclasses can override to perform
        any post-optimization processing, such as extracting solutions, running
        final measurements, or computing additional metrics.

        Args:
            **kwargs: Additional keyword arguments for subclasses.

        Note:
            The default implementation does nothing. Subclasses should override
            this method if they need post-optimization processing.
        """
        pass

    def run(self, data_file=None, **kwargs):
        """Run the variational quantum algorithm.

        The outputs are stored in the algorithm object.
        Optionally, the data can be stored in a file.

        Args:
            data_file (str | None): The file to store the data in. If None, no data is stored. Defaults to None.
            **kwargs: Additional keyword arguments for subclasses.
        """

        def cost_fn(params):
            self.reporter.info(
                message="💸 Computing Cost 💸", iteration=self.current_iteration
            )

            self._curr_params = np.atleast_2d(params)

            losses = self._run_optimization_circuits(data_file, **kwargs)

            losses = np.fromiter(losses.values(), dtype=np.float64)

            if params.ndim > 1:
                return losses
            else:
                return losses.item()

        self._grad_shift_mask = _compute_parameter_shift_mask(
            self.n_layers * self.n_params
        )

        def grad_fn(params):
            self._grad_mode = True

            self.reporter.info(
                message="📈 Computing Gradients 📈", iteration=self.current_iteration
            )

            self._curr_params = self._grad_shift_mask + params

            exp_vals = self._run_optimization_circuits(data_file, **kwargs)
            exp_vals_arr = np.fromiter(exp_vals.values(), dtype=np.float64)

            pos_shifts = exp_vals_arr[::2]
            neg_shifts = exp_vals_arr[1::2]
            grads = 0.5 * (pos_shifts - neg_shifts)

            self._grad_mode = False

            return grads

        def _iteration_counter(intermediate_result: OptimizeResult):

            self._losses_history.append(
                dict(
                    zip(
                        range(len(intermediate_result.x)),
                        intermediate_result.fun,
                    )
                )
            )

            current_loss = np.min(intermediate_result.fun)
            if current_loss < self._best_loss:
                self._best_loss = current_loss
                best_idx = np.argmin(intermediate_result.fun)

                self._best_params = intermediate_result.x[best_idx].copy()

            self.current_iteration += 1

            self.reporter.update(iteration=self.current_iteration)

            if self._cancellation_event and self._cancellation_event.is_set():
                raise _CancelledError("Cancellation requested by batch.")

            # The scipy implementation of COBYLA interprets the `maxiter` option
            # as the maximum number of function evaluations, not iterations.
            # To provide a consistent user experience, we disable `scipy`'s
            # `maxiter` and manually stop the optimization from the callback
            # when the desired number of iterations is reached.
            if (
                isinstance(self.optimizer, ScipyOptimizer)
                and self.optimizer.method == ScipyMethod.COBYLA
                and intermediate_result.nit + 1 == self.max_iterations
            ):
                raise StopIteration

        self.reporter.info(message="Finished Setup")

        # Only initialize if user hasn't already set initial_params
        if self._curr_params is None:
            self._initialize_params()

        try:
            self._minimize_res = self.optimizer.optimize(
                cost_fn=cost_fn,
                initial_params=self._curr_params,
                callback_fn=_iteration_counter,
                jac=grad_fn,
                maxiter=self.max_iterations,
                rng=self._rng,
            )
        except _CancelledError:
            # The optimizer was stopped by our callback. This is not a real
            # error, just a signal to exit this task cleanly.
            return self._total_circuit_count, self._total_run_time

        self._final_params = self._minimize_res.x

        self._perform_final_computation(**kwargs)

        self.reporter.info(message="Finished successfully!")

        return self.total_circuit_count, self.total_run_time

    def _run_solution_measurement(self):
        """Execute measurement circuits to obtain probability distributions for solution extraction."""
        if self._best_params is None:
            raise RuntimeError(
                "Optimization has not been run, no best parameters available."
            )

        if "meas_circuit" not in self._meta_circuits:
            raise NotImplementedError(
                f"{type(self).__name__} does not implement a 'meas_circuit'."
            )

        self._is_compute_probabilites = True

        # Compute probabilities for best parameters (the ones that achieved best loss)
        self._curr_params = np.atleast_2d(self._best_params)
        self._curr_circuits = self._generate_circuits()
        best_probs = self._dispatch_circuits_and_process_results()
        self._best_probs.update(best_probs)

        self._is_compute_probabilites = False
