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

import logging
import pickle
import re
from abc import abstractmethod
from datetime import datetime
from functools import partial
from pathlib import Path
from queue import Queue
from typing import Any, NamedTuple
from warnings import warn

import numpy as np
import numpy.typing as npt
import pennylane as qml
from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator
from scipy.optimize import OptimizeResult

from divi.backends import (
    CircuitRunner,
    convert_counts_to_probs,
    reverse_dict_endianness,
)
from divi.circuits import CircuitBundle, CircuitTag, MetaCircuit, format_circuit_tag
from divi.circuits.qem import _NoMitigation
from divi.qprog._expectation import _batched_expectation
from divi.qprog._hamiltonians import convert_hamiltonian_to_pauli_string
from divi.qprog.checkpointing import (
    PROGRAM_STATE_FILE,
    CheckpointConfig,
    _atomic_write,
    _ensure_checkpoint_dir,
    _get_checkpoint_subdir_path,
    _load_and_validate_pydantic_model,
    resolve_checkpoint_path,
)
from divi.qprog.exceptions import _CancelledError
from divi.qprog.optimizers import (
    MonteCarloOptimizer,
    Optimizer,
    PymooOptimizer,
    ScipyMethod,
    ScipyOptimizer,
)
from divi.qprog.quantum_program import QuantumProgram

logger = logging.getLogger(__name__)


def _get_run_instruction() -> str:
    return "Call run() to execute the optimization."


def _merge_param_group_counts(
    param_group: list[tuple[str, dict[str, int]]],
) -> dict[str, int]:
    """Merge shot histograms for a single parameter group."""
    shots_dict: dict[str, int] = {}
    for _, d in param_group:
        for s, c in d.items():
            shots_dict[s] = shots_dict.get(s, 0) + c
    return shots_dict


def _merge_probability_histograms(
    grouped: dict[
        int,
        dict[int, dict[tuple[str, int], list[dict[str, int]]]],
    ],
    shots: int,
) -> dict[str, dict[str, float]]:
    """Merge probability histograms across Hamiltonian samples per (param_id, qem_name, qem_id).

    Preserves (qem_name, qem_id) identity. Averages probabilities across hamiltonian_id
    within each QEM group instead of aggregating counts across qem_ids (which would
    over-count shots for protocols like ZNE with multiple fold factors).
    """

    merged: dict[str, dict[str, float]] = {}

    for param_id, ham_dict in grouped.items():
        qem_keys = set()
        for qem_dict in ham_dict.values():
            qem_keys.update(qem_dict.keys())

        for qem_key in sorted(qem_keys, key=lambda k: k[1]):
            qem_name, qem_id = qem_key
            probs_per_ham: list[dict[str, float]] = []

            for _, qem_dict in ham_dict.items():
                if qem_key not in qem_dict:
                    continue
                shots_list = qem_dict[qem_key]
                param_group = [("", s) for s in shots_list]
                merged_counts = _merge_param_group_counts(param_group)
                probs_per_ham.append(
                    {
                        bitstring: count / shots
                        for bitstring, count in merged_counts.items()
                    }
                )

            if not probs_per_ham:
                continue

            all_bitstrings = set()
            for p in probs_per_ham:
                all_bitstrings.update(p.keys())

            avg_probs: dict[str, float] = {}
            for bs in all_bitstrings:
                vals = [p.get(bs, 0.0) for p in probs_per_ham]
                avg_probs[bs] = sum(vals) / len(vals)

            tag = CircuitTag(
                param_id=param_id,
                qem_name=qem_name,
                qem_id=qem_id,
                hamiltonian_id=0,
                meas_id=0,
            )
            merged[format_circuit_tag(tag)] = avg_probs

    return reverse_dict_endianness(merged)


class SolutionEntry(NamedTuple):
    """A solution entry with bitstring, probability, and optional decoded value.

    Args:
        bitstring: Binary string representing a computational basis state.
        prob: Measured probability in range [0.0, 1.0].
        decoded: Optional problem-specific decoded representation. Defaults to None.
    """

    bitstring: str
    prob: float
    decoded: Any | None = None


class SubclassState(BaseModel):
    """Container for subclass-specific state."""

    data: dict[str, Any] = Field(default_factory=dict)


class OptimizerConfig(BaseModel):
    """Configuration for reconstructing an optimizer."""

    type: str
    config: dict[str, Any] = Field(default_factory=dict)


class ProgramState(BaseModel):
    """Pydantic model for VariationalQuantumAlgorithm state."""

    model_config = ConfigDict(from_attributes=True, populate_by_name=True)

    # Metadata
    program_type: str = Field(validation_alias="_serialized_program_type")
    version: str = "1.0"
    timestamp: str = Field(default_factory=lambda: datetime.now().isoformat())

    # Core Algorithm State (mapped to private attributes)
    current_iteration: int
    max_iterations: int
    losses_history: list[dict[str, float]] = Field(validation_alias="_losses_history")
    best_loss: float = Field(validation_alias="_best_loss")
    best_probs: dict[str, float] = Field(validation_alias="_best_probs")
    total_circuit_count: int = Field(validation_alias="_total_circuit_count")
    total_run_time: float = Field(validation_alias="_total_run_time")
    seed: int | None = Field(validation_alias="_seed")
    grouping_strategy: str = Field(validation_alias="_grouping_strategy")

    # Arrays
    curr_params: list[list[float]] | None = Field(
        default=None, validation_alias="_curr_params"
    )
    best_params: list[float] | None = Field(
        default=None, validation_alias="_best_params"
    )
    final_params: list[float] | None = Field(
        default=None, validation_alias="_final_params"
    )

    # Complex State (mapped to new adapter properties)
    rng_state_bytes: bytes | None = Field(
        default=None, validation_alias="_serialized_rng_state"
    )
    optimizer_config: OptimizerConfig = Field(
        validation_alias="_serialized_optimizer_config"
    )
    subclass_state: SubclassState = Field(validation_alias="_serialized_subclass_state")

    @field_serializer("rng_state_bytes")
    def serialize_bytes(self, v: bytes | None, _info):
        return v.hex() if v is not None else None

    @field_validator("rng_state_bytes", mode="before")
    @classmethod
    def validate_bytes(cls, v):
        return bytes.fromhex(v) if isinstance(v, str) else v

    @field_serializer("curr_params", "best_params", "final_params")
    def serialize_arrays(self, v: npt.NDArray | list | None, _info):
        if isinstance(v, np.ndarray):
            return v.tolist()
        return v

    def restore(self, program: "VariationalQuantumAlgorithm") -> None:
        """Apply this state object back to a program instance."""
        # 1. Bulk restore standard attributes
        for name, field in self.model_fields.items():
            target_attr = field.validation_alias or name

            # Skip adapter properties (they are read-only / calculated)
            if target_attr.startswith("_serialized_"):
                continue

            val = getattr(self, name)

            # Handle numpy conversion
            if "params" in target_attr and val is not None:
                val = np.array(val)

            if hasattr(program, target_attr):
                setattr(program, target_attr, val)

        # 2. Restore complex state
        if self.rng_state_bytes:
            program._rng.bit_generator.state = pickle.loads(self.rng_state_bytes)

        program._load_subclass_state(self.subclass_state.data)


def _compute_parameter_shift_mask(n_params: int) -> npt.NDArray[np.float64]:
    """
    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:
        npt.NDArray[np.float64]: 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 (npt.NDArray[np.float64]): Final optimized parameters.
        _best_params (npt.NDArray[np.float64]): 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 (npt.NDArray[np.float64]): 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,
        optimizer: Optimizer | None = None,
        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.
            optimizer (Optimizer | None): The optimizer to use for parameter optimization.
                Defaults to MonteCarloOptimizer().
            seed (int | None): Random seed for parameter initialization. Defaults to None.
            progress_queue (Queue | None): Queue for progress reporting. Defaults to None.

        Keyword Args:
            initial_params (npt.NDArray[np.float64] | None): Initial parameters with shape
                (n_param_sets, n_layers * n_params). If provided, these will be set as
                the current parameters via the `curr_params` setter (which includes validation).
                Defaults to None.
            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.
            precision (int): Number of decimal places for parameter values in QASM conversion.
                Defaults to 8.

                Note: Higher precision values result in longer QASM strings, which increases
                the amount of data sent to cloud backends. For most use cases, the default
                precision of 8 decimal places provides sufficient accuracy while keeping
                QASM sizes manageable. Consider reducing precision if you need to minimize
                data transfer overhead, or increase it only if you require higher numerical
                precision in your circuit parameters.
            decode_solution_fn (callable[[str], Any] | None): Function to decode bitstrings
                into problem-specific solution representations. Called during final computation
                and when `get_top_solutions(include_decoded=True)` is used. The function should
                take a binary string (e.g., "0101") and return a decoded representation
                (e.g., a list of indices, numpy array, or custom object). Defaults to
                `lambda bitstring: bitstring` (identity function).
        """

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

        # --- Optimization Results & History ---
        self._losses_history = []
        self._best_params = []
        self._final_params = []
        self._best_loss = float("inf")
        self._best_probs = {}
        self._curr_params = kwargs.pop("initial_params", None)

        # --- Random Number Generation ---
        self._seed = seed
        self._rng = np.random.default_rng(self._seed)

        # --- Computation Mode Flags ---
        # Lets child classes adapt their optimization step for grad calculation routine
        self._grad_mode = False
        self._is_compute_probabilities = False

        # --- Optimizer Configuration ---
        self.optimizer = optimizer if optimizer is not None else MonteCarloOptimizer()

        # --- Backend & Circuit Configuration ---
        if backend and backend.supports_expval:
            grouping_strategy = kwargs.pop("grouping_strategy", None)
            if grouping_strategy is not None and grouping_strategy != "_backend_expval":
                warn(
                    "Backend supports direct expectation value calculation, but a grouping_strategy was provided. "
                    "The grouping strategy will be ignored.",
                    UserWarning,
                )
            self._grouping_strategy = "_backend_expval"
        else:
            self._grouping_strategy = kwargs.pop("grouping_strategy", "qwc")

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

        # --- Solution Decoding ---
        self._decode_solution_fn = kwargs.pop(
            "decode_solution_fn", lambda bitstring: bitstring
        )

        # --- Circuit Factory & Templates ---
        self._meta_circuits = None
        self._meta_circuit_factory = partial(
            MetaCircuit,
            # No grouping strategy for expectation value measurements
            grouping_strategy=self._grouping_strategy,
            qem_protocol=self._qem_protocol,
            precision=self._precision,
        )

        # --- Control Flow ---
        self._cancellation_event = None

    @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) -> dict[str, MetaCircuit]:
        """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

    def _has_run_optimization(self) -> bool:
        """Check if optimization has been run at least once.

        Returns:
            bool: True if optimization has been run, False otherwise.
        """
        return len(self._losses_history) > 0

    @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.
        """
        if not self._has_run_optimization():
            warn(
                "losses_history is empty. Optimization has not been run yet. "
                f"{_get_run_instruction()}",
                UserWarning,
                stacklevel=2,
            )
        return self._losses_history.copy()

    @property
    def min_losses_per_iteration(self) -> list[float]:
        """Get the minimum loss value for each iteration.

        Returns a list where each element is the minimum (best) loss value
        across all parameter sets for that iteration.

        Returns:
            list[float]: List of minimum loss values, one per iteration.
        """
        if not self._has_run_optimization():
            warn(
                "min_losses_per_iteration is empty. Optimization has not been run yet. "
                f"{_get_run_instruction()}",
                UserWarning,
                stacklevel=2,
            )
        return [min(loss_dict.values()) for loss_dict in self._losses_history]

    @property
    def final_params(self) -> npt.NDArray[np.float64]:
        """Get a copy of the final optimized parameters.

        Returns:
            npt.NDArray[np.float64]: Copy of the final parameters. Modifications to this array
                will not affect the internal state.
        """
        if len(self._final_params) == 0 or not self._has_run_optimization():
            warn(
                "final_params is not available. Optimization has not been run yet. "
                f"{_get_run_instruction()}",
                UserWarning,
                stacklevel=2,
            )
        return self._final_params.copy()

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

        Returns:
            npt.NDArray[np.float64]: Copy of the best parameters. Modifications to this array
                will not affect the internal state.
        """
        if len(self._best_params) == 0 or not self._has_run_optimization():
            warn(
                "best_params is not available. Optimization has not been run yet. "
                f"{_get_run_instruction()}",
                UserWarning,
                stacklevel=2,
            )
        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.
        """
        if not self._has_run_optimization():
            warn(
                "best_loss has not been computed yet. Optimization has not been run. "
                f"{_get_run_instruction()}",
                UserWarning,
                stacklevel=2,
            )
        elif self._best_loss == float("inf"):
            # Defensive check: if optimization ran but best_loss is still inf, something is wrong
            raise RuntimeError(
                "best_loss is still infinite after optimization. This indicates a problem "
                "with the optimization process. The optimization callback may not have executed "
                "correctly, or all computed losses were infinite."
            )
        return self._best_loss

    @property
    def best_probs(self) -> dict[CircuitTag, dict[str, float]]:
        """Get normalized probabilities for the best parameters.

        This property provides access to the probability distribution computed
        by running measurement circuits with the best parameters found during
        optimization. The distribution maps bitstrings (computational basis states)
        to their measured probabilities.

        The probabilities are normalized and have deterministic ordering when
        iterated (dictionary insertion order is preserved in Python 3.7+).

        Returns:
            dict[CircuitTag, dict[str, float]]: Dictionary mapping CircuitTag keys to
                bitstring probability dictionaries. Bitstrings are binary strings
                (e.g., "0101"), values are probabilities in range [0.0, 1.0].
                Returns an empty dict if final computation has not been performed.

        Raises:
            RuntimeError: If attempting to access probabilities before running
                the algorithm with final computation enabled.

        Note:
            To populate this distribution, you must run the algorithm with
            `perform_final_computation=True` (the default):

            >>> program.run(perform_final_computation=True)
            >>> probs = program.best_probs

        Example:
            >>> program.run()
            >>> probs = program.best_probs
            >>> for bitstring, prob in probs.items():
            ...     print(f"{bitstring}: {prob:.2%}")
            0101: 42.50%
            1010: 31.20%
            ...
        """
        if not self._best_probs:
            warn(
                "best_probs is empty. Either optimization has not been run yet, "
                f"or final computation was not performed. {_get_run_instruction()}",
                UserWarning,
                stacklevel=2,
            )
        return self._best_probs.copy()

    def get_top_solutions(
        self, n: int = 10, *, min_prob: float = 0.0, include_decoded: bool = False
    ) -> list[SolutionEntry]:
        """Get the top-N solutions sorted by probability.

        This method extracts the most probable solutions from the measured
        probability distribution. Solutions are sorted by probability (descending)
        with deterministic tie-breaking using lexicographic ordering of bitstrings.

        Args:
            n (int): Maximum number of solutions to return. Must be non-negative.
                If n is 0 or negative, returns an empty list. If n exceeds the
                number of available solutions (after filtering), returns all
                available solutions. Defaults to 10.
            min_prob (float): Minimum probability threshold for including solutions.
                Only solutions with probability >= min_prob will be included.
                Must be in range [0.0, 1.0]. Defaults to 0.0 (no filtering).
            include_decoded (bool): Whether to populate the `decoded` field of
                each SolutionEntry by calling the `decode_solution_fn` provided
                in the constructor. If False, the decoded field will be None.
                Defaults to False.

        Returns:
            list[SolutionEntry]: List of solution entries sorted by probability
                (descending), then by bitstring (lexicographically ascending)
                for deterministic tie-breaking. Returns an empty list if no
                probability distribution is available or n <= 0.

        Raises:
            RuntimeError: If probability distribution is not available because
                optimization has not been run or final computation was not performed.
            ValueError: If min_prob is not in range [0.0, 1.0] or n is negative.

        Note:
            The probability distribution must be computed by running the algorithm
            with `perform_final_computation=True` (the default):

            >>> program.run(perform_final_computation=True)
            >>> top_10 = program.get_top_solutions(n=10)

        Example:
            >>> # Get top 5 solutions with probability >= 5%
            >>> program.run()
            >>> solutions = program.get_top_solutions(n=5, min_prob=0.05)
            >>> for sol in solutions:
            ...     print(f"{sol.bitstring}: {sol.prob:.2%}")
            1010: 42.50%
            0101: 31.20%
            1100: 15.30%
            0011: 8.50%
            1111: 2.50%

            >>> # Get solutions with decoding
            >>> solutions = program.get_top_solutions(n=3, include_decoded=True)
            >>> for sol in solutions:
            ...     print(f"{sol.bitstring} -> {sol.decoded}")
            1010 -> [0, 2]
            0101 -> [1, 3]
            ...
        """
        # Validate inputs
        if n < 0:
            raise ValueError(f"n must be non-negative, got {n}")
        if not (0.0 <= min_prob <= 1.0):
            raise ValueError(f"min_prob must be in range [0.0, 1.0], got {min_prob}")

        # Handle edge case: n == 0
        if n == 0:
            return []

        # Require probability distribution to exist
        if not self._best_probs:
            raise RuntimeError(
                "No probability distribution available. The final computation step "
                "must be performed to compute the probability distribution. "
                "Call run(perform_final_computation=True) to execute optimization "
                "and compute the distribution."
            )
        # Extract the probability distribution (nested by parameter set)
        # _best_probs structure: {tag: {bitstring: prob}}
        probs_dict = next(iter(self._best_probs.values()))

        # Filter by minimum probability and get top n sorted by probability (descending),
        # then bitstring (ascending) for deterministic tie-breaking
        top_items = sorted(
            filter(
                lambda bitstring_prob: bitstring_prob[1] >= min_prob, probs_dict.items()
            ),
            key=lambda bitstring_prob: (-bitstring_prob[1], bitstring_prob[0]),
        )[:n]

        # Build result list (decode on demand)
        return [
            SolutionEntry(
                bitstring=bitstring,
                prob=prob,
                decoded=(
                    self._decode_solution_fn(bitstring) if include_decoded else None
                ),
            )
            for bitstring, prob in top_items
        ]

    @property
    def curr_params(self) -> npt.NDArray[np.float64]:
        """Get the current parameters.

        These are the parameters used for optimization. They can be accessed
        and modified at any time, including during optimization.

        Returns:
            npt.NDArray[np.float64]: Current parameters. If not yet initialized,
                they will be generated automatically.
        """
        if self._curr_params is None:
            self._initialize_params()
        return self._curr_params.copy()

    @curr_params.setter
    def curr_params(self, value: npt.NDArray[np.float64] | None):
        """
        Set the current parameters.

        Args:
            value (npt.NDArray[np.float64] | None): 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

    # --- Serialization Adapters (For Pydantic) ---
    @property
    def _serialized_program_type(self) -> str:
        return type(self).__name__

    @property
    def _serialized_rng_state(self) -> bytes:
        return pickle.dumps(self._rng.bit_generator.state)

    @property
    def _serialized_optimizer_config(self) -> OptimizerConfig:
        config_dict = self.optimizer.get_config()
        return OptimizerConfig(type=config_dict.pop("type"), config=config_dict)

    @property
    def _serialized_subclass_state(self) -> SubclassState:
        return SubclassState(data=self._save_subclass_state())

    @property
    def meta_circuits(self) -> dict[str, MetaCircuit]:
        """Get the meta-circuit templates used by this program.

        Returns:
            dict[str, MetaCircuit]: Dictionary mapping circuit names to their
                MetaCircuit templates.
        """
        # Lazy initialization: each instance has its own _meta_circuits.
        # Note: When used with ProgramBatch, meta_circuits is initialized sequentially
        # in the main thread before parallel execution to avoid thread-safety issues.
        if self._meta_circuits is None:
            self._meta_circuits = self._create_meta_circuits_dict()
        return self._meta_circuits

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

    @abstractmethod
    def _generate_circuits(self, **kwargs) -> list[CircuitBundle]:
        """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[CircuitBundle]: List of Circuit objects to be executed.
        """
        pass

    @abstractmethod
    def _save_subclass_state(self) -> dict[str, Any]:
        """Hook method for subclasses to save additional state.

        Subclasses must override this method to return a dictionary of
        state variables that should be included in the checkpoint.

        Returns:
            dict[str, Any]: Dictionary of subclass-specific state.
        """
        pass

    @abstractmethod
    def _load_subclass_state(self, state: dict[str, Any]) -> None:
        """Hook method for subclasses to load additional state.

        Subclasses must override this method to restore state variables
        from the checkpoint dictionary. This is called after instance creation.

        Args:
            state (dict[str, Any]): Dictionary of subclass-specific state.
        """
        pass

    def _get_optimizer_config(self) -> OptimizerConfig:
        """Extract optimizer configuration for checkpoint reconstruction.

        Returns:
            OptimizerConfig: Configuration object for the current optimizer.

        Raises:
            NotImplementedError: If the optimizer does not support state saving.
        """
        config_dict = self.optimizer.get_config()
        return OptimizerConfig(
            type=config_dict.pop("type"),
            config=config_dict,
        )

    def save_state(self, checkpoint_config: CheckpointConfig) -> str:
        """Save the program state to a checkpoint directory."""
        if self.current_iteration == 0 and len(self._losses_history) == 0:
            raise RuntimeError("Cannot save checkpoint: optimization has not been run.")

        if checkpoint_config.checkpoint_dir is None:
            raise ValueError(
                "checkpoint_config.checkpoint_dir must be a non-None Path."
            )

        main_dir = _ensure_checkpoint_dir(checkpoint_config.checkpoint_dir)
        checkpoint_path = _get_checkpoint_subdir_path(main_dir, self.current_iteration)
        checkpoint_path.mkdir(parents=True, exist_ok=True)

        # 1. Save optimizer
        self.optimizer.save_state(checkpoint_path)

        # 2. Save Program State (Pydantic pulls data via validation_aliases)
        state = ProgramState.model_validate(self)

        state_file = checkpoint_path / PROGRAM_STATE_FILE
        _atomic_write(state_file, state.model_dump_json(indent=2))

        return checkpoint_path

    @classmethod
    def load_state(
        cls,
        checkpoint_dir: Path | str,
        backend: CircuitRunner,
        subdirectory: str | None = None,
        **kwargs,
    ) -> "VariationalQuantumAlgorithm":
        """Load program state from a checkpoint directory."""
        checkpoint_path = resolve_checkpoint_path(checkpoint_dir, subdirectory)
        state_file = checkpoint_path / PROGRAM_STATE_FILE

        # 1. Load Pydantic Model
        state = _load_and_validate_pydantic_model(
            state_file,
            ProgramState,
            required_fields=["program_type", "current_iteration"],
        )

        # 2. Reconstruct Optimizer
        opt_config = state.optimizer_config
        if opt_config.type == "MonteCarloOptimizer":
            optimizer = MonteCarloOptimizer.load_state(checkpoint_path)
        elif opt_config.type == "PymooOptimizer":
            optimizer = PymooOptimizer.load_state(checkpoint_path)
        else:
            raise ValueError(f"Unsupported optimizer type: {opt_config.type}")

        # 3. Create Instance
        program = cls(backend=backend, optimizer=optimizer, seed=state.seed, **kwargs)

        # 4. Restore State
        state.restore(program)

        return program

    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: npt.NDArray[np.float64]):
        """
        Validate user-provided initial parameters.

        Args:
            params (npt.NDArray[np.float64]): 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, **kwargs) -> dict[int, float]:
        self._curr_circuits = self._generate_circuits(**kwargs)

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

        losses = self._dispatch_circuits_and_process_results(**kwargs)

        return losses

    @staticmethod
    def _parse_tag_full(tag: CircuitTag | str) -> tuple[int, str, int, int, int]:
        """Extract (param_id, qem_name, qem_id, hamiltonian_id, meas_id) from a tag."""
        if isinstance(tag, CircuitTag):
            return (
                tag.param_id,
                tag.qem_name,
                tag.qem_id,
                tag.hamiltonian_id,
                tag.meas_id,
            )
        m = re.match(r"^(\d+)_([^:]+):(\d+)_ham:(-?\d+)_(\d+)$", str(tag))
        if not m:
            raise ValueError(f"Cannot parse tag: {tag!r}")
        return (
            int(m.group(1)),
            m.group(2),
            int(m.group(3)),
            int(m.group(4)),
            int(m.group(5)),
        )

    def _group_results(self, results: dict[str, dict[str, int]]) -> dict[
        int,
        dict[int, dict[tuple[str, int], list[dict[str, int]]]],
    ]:
        """
        Group results by parameter id, hamiltonian id (for multi-sample QDrift),
        and (qem_name, qem_id).

        Returns {param_id: {hamiltonian_id: {(qem_name, qem_id): [shots...]}}}.

        Shots within each (qem_name, qem_id) list are ordered by meas_id.
        """
        grouped: dict[
            int,
            dict[int, dict[tuple[str, int], list[tuple[int, dict[str, int]]]]],
        ] = {}
        for tag, shots in results.items():
            param_id, qem_name, qem_id, hamiltonian_id, meas_id = self._parse_tag_full(
                tag
            )
            qem_key = (qem_name, qem_id)
            inner = (
                grouped.setdefault(param_id, {})
                .setdefault(hamiltonian_id, {})
                .setdefault(qem_key, [])
            )
            inner.append((meas_id, shots))

        def strip_meas_id(qem_dict):
            return {
                q: [s for _, s in sorted(meas_shots, key=lambda x: x[0])]
                for q, meas_shots in qem_dict.items()
            }

        result: dict[
            int,
            dict[int, dict[tuple[str, int], list[dict[str, int]]]],
        ] = {}

        for p, ham_dict in grouped.items():
            result[p] = {h: strip_meas_id(qem_dict) for h, qem_dict in ham_dict.items()}

        return result

    def _reset_tag_cache(self) -> None:
        """Reset per-run tag cache for structured result tags."""
        self._tag_map: dict[str, CircuitTag] = {}

    def _encode_tag(self, tag: CircuitTag | str) -> str:
        """Convert structured tags to backend-safe strings."""
        if isinstance(tag, CircuitTag):
            tag_str = format_circuit_tag(tag)
            self._tag_map[tag_str] = tag
            return tag_str
        return str(tag)

    def _decode_tags(
        self, results: dict[str, dict[str, int]]
    ) -> dict[CircuitTag | str, dict[str, int]]:
        """Restore structured tags from backend result labels."""
        if not self._tag_map:
            return results
        return {self._tag_map.get(tag, tag): shots for tag, shots in results.items()}

    def _apply_qem_protocol(
        self, exp_matrix: npt.NDArray[np.float64]
    ) -> list[npt.NDArray[np.float64]]:
        """Apply the configured QEM protocol to expectation value matrices."""
        return [
            self._qem_protocol.postprocess_results(exp_vals) for exp_vals in exp_matrix
        ]

    def _compute_marginal_results(
        self,
        qem_groups: dict[tuple[str, int], list[dict[str, int]]],
        measurement_groups: list[list[qml.operation.Operator]],
        ham_ops: str | None,
    ) -> list[npt.NDArray[np.float64] | list[npt.NDArray[np.float64]]]:
        """Compute marginal results, handling backend modes and QEM."""
        qem_sorted = sorted(qem_groups.items(), key=lambda x: x[0][1])
        if self.backend.supports_expval:
            if ham_ops is None:
                raise ValueError(
                    "Hamiltonian operators (ham_ops) are required when using a backend "
                    "that supports expectation values, but were not provided."
                )
            ham_ops_list = ham_ops.split(";")
            qem_group_values = [shots for _, shots in qem_sorted]
            return [
                self._apply_qem_protocol(
                    np.array(
                        [
                            [shot_dict[op] for op in ham_ops_list]
                            for shot_dict in shots_dicts
                        ]
                    ).T
                )
                for shots_dicts in qem_group_values
            ] or []

        qem_group_values = [shots for _, shots in qem_sorted]
        shots_by_qem_idx = zip(*qem_group_values)
        marginal_results: list[
            npt.NDArray[np.float64] | list[npt.NDArray[np.float64]]
        ] = []
        wire_order = tuple(reversed(self.cost_hamiltonian.wires))
        for shots_dicts, curr_measurement_group in zip(
            shots_by_qem_idx, measurement_groups
        ):
            exp_matrix = _batched_expectation(
                shots_dicts, curr_measurement_group, wire_order
            )
            mitigated = self._apply_qem_protocol(exp_matrix)
            marginal_results.append(mitigated if len(mitigated) > 1 else mitigated[0])

        return marginal_results

    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[CircuitTag, dict[str, int]]): The shot histograms of the quantum execution
                step. Keys are CircuitTag instances containing param, QEM, and measurement ids.

        Returns:
            dict[int, float]: The energies for each parameter set grouping, where the dict keys
                correspond to the parameter indices.
        """
        if self._is_compute_probabilities:
            grouped = self._group_results(results)
            has_multi_ham = any(len(ham_dict) > 1 for ham_dict in grouped.values())
            if has_multi_ham:
                return _merge_probability_histograms(grouped, self.backend.shots)
            probs = convert_counts_to_probs(results, self.backend.shots)
            return reverse_dict_endianness(probs)

        if not (self._cancellation_event and self._cancellation_event.is_set()):
            self.reporter.info(
                message="Post-processing output", iteration=self.current_iteration
            )

        losses = {}
        ham_ops = kwargs.get("ham_ops")

        for p, ham_dict in self._group_results(results).items():
            energies: list[float] = []
            for ham_id, qem_groups in ham_dict.items():
                cost_circuit = (
                    self._get_cost_circuit_for_hamiltonian(ham_id)
                    if hasattr(self, "_get_cost_circuit_for_hamiltonian")
                    else self.meta_circuits["cost_circuit"]
                )
                measurement_groups = cost_circuit.measurement_groups
                marginal_results = self._compute_marginal_results(
                    qem_groups=qem_groups,
                    measurement_groups=measurement_groups,
                    ham_ops=ham_ops,
                )

                pl_loss = cost_circuit.postprocessing_fn(marginal_results).item()
                energies.append(pl_loss)

            losses[p] = np.mean(energies) + self.loss_constant

        return losses

    def _perform_final_computation(self, **kwargs) -> None:
        """
        Perform final computations after optimization is complete.

        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,
        perform_final_computation: bool = True,
        checkpoint_config: CheckpointConfig | None = None,
        **kwargs,
    ) -> tuple[int, float]:
        """Run the variational quantum algorithm.

        The outputs are stored in the algorithm object.

        Args:
            perform_final_computation (bool): Whether to perform final computation after optimization completes.
                Typically, this step involves sampling with the best found parameters to extract
                solution probability distributions. Set this to False in warm-starting or pre-training
                routines where the final sampling step is not needed. Defaults to True.
            checkpoint_config (CheckpointConfig | None): Checkpoint configuration.
                If None, no checkpointing is performed.
            **kwargs: Additional keyword arguments for subclasses.

        Returns:
            tuple[int, float]: A tuple containing (total_circuit_count, total_run_time).
        """
        # Initialize checkpointing
        if checkpoint_config is None:
            checkpoint_config = CheckpointConfig()

        if checkpoint_config.checkpoint_dir:
            logger.info(
                f"Using checkpoint directory: {checkpoint_config.checkpoint_dir}"
            )

        # Extract max_iterations from kwargs if present (for compatibility with subclasses)
        max_iterations = kwargs.pop("max_iterations", self.max_iterations)
        if max_iterations != self.max_iterations:
            self.max_iterations = max_iterations

        # Warn if max_iterations is less than current_iteration (regardless of how it was set)
        if self.max_iterations < self.current_iteration:
            warn(
                f"max_iterations ({self.max_iterations}) is less than current_iteration "
                f"({self.current_iteration}). The optimization will not run additional "
                f"iterations since the maximum has already been reached.",
                UserWarning,
            )

        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(**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(**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(
                        [str(i) for i in 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)

            # Checkpointing
            if checkpoint_config._should_checkpoint(self.current_iteration):
                self.save_state(checkpoint_config)

            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")

        if self._curr_params is None:
            self._initialize_params()
        else:
            self._validate_initial_params(self._curr_params)

        try:
            self._minimize_res = self.optimizer.optimize(
                cost_fn=cost_fn,
                initial_params=self._curr_params,
                callback_fn=_iteration_counter,
                jac=grad_fn,
                max_iterations=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

        # Set _best_params from final result (source of truth)
        x = np.atleast_2d(self._minimize_res.x)
        fun = np.atleast_1d(self._minimize_res.fun)
        best_idx = np.argmin(fun)
        self._best_params = x[best_idx].copy()

        if perform_final_computation:
            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) -> None:
        """Execute measurement circuits to obtain probability distributions for solution extraction."""

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

        self._is_compute_probabilities = 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_probabilities = False
