import structlog

LOG = structlog.get_logger(file=__name__)

import numpy as np

from qiskit import (
    QuantumCircuit,
    QuantumRegister,
    ClassicalRegister,
    AncillaRegister,
    transpile,
)
from qiskit.circuit import Parameter
try:
    from qiskit_aer.primitives import SamplerV2 as _SamplerV2
except ImportError:
    from qiskit.primitives import StatevectorSampler as _SamplerV2
from qiskit_algorithms.optimizers import COBYLA

from qiskit_aer import AerSimulator, Aer

from qaoa.initialstates import InitialState
from qaoa.mixers import Mixer
from qaoa.problems import Problem
from qaoa.util import Statistic, BitFlip, post_processing


class OptResult:
    """
    Class for results of optimization at depth p.

    Attributes:
        depth (int): The depth p of the optimization.
        angles (list): List of angles used in the optimization.
        Exp (list): List of expected values.
        Var (list): List of variances.
        WorstCost (list): List of worst costs.
        BestCost (list): List of best costs.
        BestSols (list): List of best solutions.
        shots (list): List of shots taken for each iteration.
        index_Exp_min (int): Index of the minimum expected value.

    Methods:
        add_iteration(): Adds an iteration's results to the OptResult object.
        compute_best_index(): Computes the index of the minimum expected value in `Exp`.
        get_best_Exp(): Returns the best expected value.
        get_best_Var(): Returns the best variance.
        get_best_angles(): Returns the best angles corresponding to the minimum expected value.
        num_fval(): Returns the number of function evaluations.
        num_shots(): Returns the total number of shots taken.
        get_best_solution(): Returns the best solutions and their corresponding cost.
    """

    def __init__(self, depth):
        """
        Initializes the OptResult object with the given depth.

        Args:
            depth (int): The depth p of the optimization.
        """
        self.depth = depth

        self.angles = []
        self.Exp = []
        self.Var = []
        self.WorstCost = []
        self.BestCost = []
        self.BestSols = []
        self.shots = []

        self.index_Exp_min = -1

    def add_iteration(self, angles, stat, shots):
        """
        Adds an iteration's results to the OptResult object.

        Args:
            angles (list): List of angles used in the iteration.
            stat (Statistic): Statistic object containing the results of the iteration.
            shots (int): Number of shots taken in the iteration.
        """
        self.angles.append(angles)
        self.Exp.append(-stat.get_CVaR())
        self.Var.append(stat.get_Variance())
        self.BestCost.append(-stat.get_max())
        self.WorstCost.append(-stat.get_min())
        self.BestSols.append(stat.get_max_sols())
        self.shots.append(shots)

    def compute_best_index(self):
        """
        Computes the index of the minimum expected value in `Exp`.
        This index is used to retrieve the best angles, expected value, and variance.
        """
        self.index_Exp_min = self.Exp.index(min(self.Exp))

    def get_best_Exp(self):
        """
        Returns:
            float: The best expected value from `Exp`.
        """
        return self.Exp[self.index_Exp_min]

    def get_best_Var(self):
        """
        Returns:
            float: The best variance from `Var`.
        """
        return self.Var[self.index_Exp_min]

    def get_best_angles(self):
        """
        Returns:
            list: The best angles from `angles` at the index of the minimum expected value.
        """
        return self.angles[self.index_Exp_min]

    def num_fval(self):
        """
        Returns:
            int: The number of function evaluations (the length of `Exp`).
        """
        return len(self.Exp)

    def num_shots(self):
        """
        Returns:
            int: The total number of shots taken across all iterations.
        """
        return sum(self.shots)

    def get_best_solution(self):
        """
        Iterates through the best solutions and returns the best cost and the corresponding solutions.

        Returns:
            tuple: A tuple containing
                - list: The best solutions (bit-strings) that yield the best cost.
                - float: The best cost found.
        """
        best_cost = np.min(self.BestCost)
        iterations_with_best_cost = np.where(self.BestCost == best_cost)[0]

        all_best_sols = []
        for i in iterations_with_best_cost:
            all_best_sols.append(self.BestSols[i])
        # flatten the list:
        all_best_sols = [item for sublist in all_best_sols for item in sublist]
        best_sols = np.unique(all_best_sols)
        return best_sols, best_cost


class QAOA:
    """
    Quantum Approximate Optimization Algorithm (QAOA) class.

    This class implements the QAOA algorithm for solving combinatorial optimization problems
    using quantum circuits.

    The class integrates a **problem**, **mixer**, and **initial state** to create a custom ansatz.
    The library already contains standard implementations.

    Attributes:
        problem (Problem): The optimization problem to solve.
        mixer (Mixer): The mixer circuit used in the QAOA algorithm.
        initialstate (InitialState): The initial state of the quantum circuit.

        backend (AerBackend): The backend to run the quantum circuits on. By default it uses quiskit's QuasmSimulator.
        noisemodel (optional): Optional Qiskit noise model.
        optimizer (list): [Optimizer, {options}]. Default is COBYLA with no options.
        precision (float, optional): Optional target precision for expectation value estimation.
        shots (int): Number of measurement shots per circuit evaluation.
        cvar (float): CVaR parameter for expectation value calculation.
        memorysize (int): Number of measurement outcomes to store (if > 0).
        interpolate (bool): Whether to use parameter interpolation for initial guesses.
        flip (bool): Whether to apply bit-flip boosting between layers.
        post (bool): Whether to apply post-processing to measurement results.
        number_trottersteps_mixer (int): Number of Trotter steps for the mixer.
        sequential (bool): Whether to run sequential circuit evaluations.

        parameterized_circuit (QuantumCircuit): The parameterized quantum circuit for QAOA.
        optimization_results (dict[int, OptResult]): Dictionary mapping depth (int) to corresponding optimization result objects (OptResult).

    Main methods:
        - createParameterizedCircuit(): Creates the parameterized circuit for a given depth.
        - sample_cost_landscape(): Samples the cost landscape for given angles at a specific depth.
        - optimize(): Runs the optimization process for a given depth.
        - get_Exp(): Returns the best expected value at a given depth.
        - get_Var(): Returns the best variance at a given depth.
        - get_beta(): Returns the beta angles at a given depth.
        - get_gamma(): Returns the gamma angles at a given depth.

    """

    def __init__(
        self,
        problem,
        mixer,
        initialstate,
        backend=AerSimulator(),
        noisemodel=None,
        optimizer=[COBYLA, {}],  # optimizer, options
        precision=None,
        shots=1024,
        cvar=1,
        memorysize=-1,
        interpolate=True,
        flip=False,
        post=False,
        number_trottersteps_mixer=1,
        sequential=False,
    ) -> None:
        """
        Initializes the QAOA class.

        Args:
            problem (Problem): The optimization problem to solve.
            mixer (Mixer): The mixer circuit used in the QAOA algorithm.
            initialstate (InitialState): The initial state of the quantum circuit.

            backend (AerBackend): The backend to run the quantum circuits on. Default is Quiskit's `qasm_simulator`.
            noisemodel (optional): Optional Qiskit noise model. Default is None.
            optimizer (list): [Optimizer, {options}]. Default is COBYLA with no options.
            precision (float, optional): Optional target precision for expectation value estimation. Default is None.
            shots (int): Number of measurement shots per circuit evaluation. Default is 1024.
            cvar (float): CVaR parameter for expectation value calculation. Default is 1.
            memorysize (int): Number of measurement outcomes to store (if > 0). Default is -1 (no memory).
            interpolate (bool): Whether to use parameter interpolation for initial guesses. Default is True.
            flip (bool): Whether to apply bit-flip boosting between layers. Default is False.
            post (bool): Whether to apply post-processing to measurement results. Default is False.
            number_trottersteps_mixer (int): Number of Trotter steps for the mixer. Default is 1.
            sequential (bool): Whether to run sequential circuit evaluations. Default is False.

        """
        # A QAO-Ansatz consist of these parts:

        # :problem of Basetype Problem,
        # implementing the phase circuit and the cost.

        # :mixer of Basetype Mixer,
        # specifying the mixer circuit.

        # :initialstate of Basetype InitialState,
        # specifying the initial state.

        # :backend: backend
        # :noisemodel: noisemodel
        # :optmizer: optimizer as a list containing [optimizer object, options dict]
        # :precision: precision to reach for expectation value based on error=variance/sqrt(shots)
        # :shots: if precision=None, the number of samples taken
        #               if precision!=None, the minimum number of samples taken
        # :cvar: used for CVar
        # """

        assert issubclass(type(problem), Problem)
        assert issubclass(type(mixer), Mixer)
        assert issubclass(type(initialstate), InitialState)

        self.problem = problem
        self.mixer = mixer
        self.initialstate = initialstate
        self.initialstate.setNumQubits(self.problem.N_qubits)
        self.mixer.setNumQubits(self.problem.N_qubits)

        self.parameterized_circuit = None
        self.parameterized_circuit_depth = 0

        self.backend = backend
        self.optimizer = optimizer
        self.noisemodel = noisemodel
        self.shots = shots
        self.precision = precision
        self.stat = Statistic(cvar=cvar)
        self.cvar = cvar
        self.memorysize = memorysize
        self.memory = self.memorysize > 0
        self.interpolate = interpolate

        self.sequential = sequential

        self.usebarrier = False
        self.isQNSPSA = False

        self.current_depth = 0  # depth at which local optimization has been done

        self.gamma_params = None
        self.beta_params = None
        self.init_params = None
        self.n_gamma = 1
        self.n_beta = 1
        self.n_init = 0

        self.Exp_sampled_p1 = None
        self.landscape_p1_angles = {}
        self.Var_sampled_p1 = None
        self.MaxCost_sampled_p1 = None
        self.MinCost_sampled_p1 = None

        self.optimization_results = {}
        self.memory_lists = []

        self.flip = flip
        self.flipper = BitFlip(self.problem.N_qubits)
        self.bitflips = []

        self.post = post
        self.Exp_post_processed = None
        self.Var_post_processed = None
        self.samplecount_hists = {}
        self.last_hist = {}
        self.number_trottersteps_mixer = number_trottersteps_mixer

    def exp_landscape(self):
        """
        Returns:
            float: The expected value of the cost landscape at depth p = 1.
        """
        ### at depth p = 1
        return self.Exp_sampled_p1

    def var_landscape(self):
        """
        Returns:
            float: The variance of the cost landscape at depth p = 1.
        """
        ### at depth p = 1
        return self.Var_sampled_p1

    def get_Exp(self, depth=None):
        """
        Args:
            depth (int, optional): The depth at which to retrieve the expected value.

        Returns:
            list/float: The best expected value(s) at the specified depth.
            If depth is None, returns a list of best expected values for all depths up to the current depth.
            If depth is specified, returns the best expected value at that depth.
        """
        if not depth:
            ret = []
            for i in range(1, self.current_depth + 1):
                ret.append(self.optimization_results[i].get_best_Exp())
            return ret
        if depth > self.current_depth + 1:
            raise ValueError
        return self.optimization_results[depth].get_best_Exp()

    def get_Var(self, depth):
        """
        Args:
            depth (int): The depth at which to retrieve the variance.

        Raises:
            ValueError: If the specified depth is greater than the current depth + 1.

        Returns:
            float: The best variance at the specified depth.
        """
        if depth > self.current_depth + 1:
            raise ValueError
        return self.optimization_results[depth].get_best_Var()

    def get_beta(self, depth):
        """
        Args:
            depth (int): The depth at which to retrieve the beta angles.

        Raises:
            ValueError: If the specified depth is greater than the current depth + 1.

        Returns:
            list: The best beta angles at the specified depth.
        """
        if depth > self.current_depth + 1:
            raise ValueError
        return self.optimization_results[depth].get_best_angles()[1::2]

    def get_gamma(self, depth):
        """
        Args:
            depth (int): The depth at which to retrieve the beta angles.

        Raises:
            ValueError: If the specified depth is greater than the current depth + 1.

        Returns:
            list: The best gamma angles at the specified depth.
        """
        if depth > self.current_depth + 1:
            raise ValueError
        return self.optimization_results[depth].get_best_angles()[::2]

    def get_angles(self, depth):
        """
        Args:
            depth (int): The depth at which to retrieve the angles.
        Raises:
            ValueError: If the specified depth is greater than the current depth + 1.

        Returns:
            list: The best angles at the specified depth.
        """
        if depth > self.current_depth + 1:
            raise ValueError
        return self.optimization_results[depth].get_best_angles()

    def get_memory_lists(self):
        """
        Returns:
            list: A list of memory lists containing measurement outcomes for each depth.
        """
        return self.memory_lists

    def createParameterizedCircuit(self, depth):
        """Creates the parameterized quantum circuit for the QAOA algorithm at a specified depth.

        The function generates quantum circuits for the problem and mixer Hamiltonians and the initial state.
        The main circuit is constructed layer by layer starting from the initial state.
        For each layer, symbolic parameters are assigned to the problem unitary (gamma) and the mixer unitary (beta).
        Multi-angle components can use multiple parameters per layer.

        If `flip` is set to True, a bit-flip circuit is added between layers to modify the state according to a pre-defined bit-flip pattern.
        If `usebarrier` is set to True, visual barrier instructions are added between layers to improve circuit readability and potentially
        guide compiler optimizations.

        Finally, the full parameterized circuit is measured and transpiled for the target backend. The constructed depth is stored for reference.

        Args:
            depth (int): The depth p (number of alternating mixing and phase separating layers) of the QAOA circuit.
        """
        if self.parameterized_circuit_depth == depth:
            LOG.info(
                "Circuit is already of depth " + str(self.parameterized_circuit_depth)
            )
            return

        self.problem.create_circuit()
        self.mixer.create_circuit()
        self.initialstate.create_circuit()

        # Query number of parameters from each component
        self.n_gamma = self.problem.get_num_parameters()
        self.n_beta = self.mixer.get_num_parameters()
        self.n_init = self.initialstate.get_num_parameters()

        a = AncillaRegister(self.problem.N_ancilla_qubits)
        q = QuantumRegister(self.problem.N_qubits)
        c = ClassicalRegister(self.problem.N_qubits)
        self.parameterized_circuit = QuantumCircuit(q, c, a)

        self.gamma_params = [[None] * self.n_gamma for _ in range(depth)]
        self.beta_params = [[None] * self.n_beta for _ in range(depth)]
        self.init_params = [None] * self.n_init

        # Compose initial state circuit (with parameters if applicable)
        if self.n_init > 0:
            for i in range(self.n_init):
                self.init_params[i] = Parameter("init_{}".format(i))
            init_param_map = {
                self.initialstate.circuit.parameters[i]: self.init_params[i]
                for i in range(self.n_init)
            }
            tmp_init = self.initialstate.circuit.assign_parameters(
                init_param_map, inplace=False
            )
            self.parameterized_circuit.compose(tmp_init, inplace=True)
        else:
            self.parameterized_circuit.compose(self.initialstate.circuit, inplace=True)

        if self.usebarrier:
            self.parameterized_circuit.barrier()

        for d in range(depth):
            for i in range(self.n_gamma):
                self.gamma_params[d][i] = Parameter("gamma_{}_{}".format(d, i))
            gamma_param_map = {
                self.problem.circuit.parameters[i]: self.gamma_params[d][i]
                for i in range(self.n_gamma)
            }
            tmp_circuit = self.problem.circuit.assign_parameters(
                gamma_param_map, inplace=False
            )
            self.parameterized_circuit.compose(tmp_circuit, inplace=True)

            if self.usebarrier:
                self.parameterized_circuit.barrier()

            for i in range(self.n_beta):
                self.beta_params[d][i] = Parameter("beta_{}_{}".format(d, i))
            beta_param_map = {
                self.mixer.circuit.parameters[i]: self.beta_params[d][i]
                / self.number_trottersteps_mixer
                for i in range(self.n_beta)
            }
            tmp_circuit = self.mixer.circuit.assign_parameters(
                beta_param_map, inplace=False
            )
            for _ in range(0, self.number_trottersteps_mixer):
                self.parameterized_circuit.compose(tmp_circuit, inplace=True)

            if self.flip and (d != (depth - 1)):
                self.parameterized_circuit.barrier()
                self.flipper.create_circuit(self.bitflips[d])
                self.parameterized_circuit.compose(self.flipper.circuit, inplace=True)
                self.parameterized_circuit.barrier()

            if self.usebarrier:
                self.circuit.barrier()

        self.parameterized_circuit.barrier()
        self.parameterized_circuit.measure(q, c)
        self.parameterized_circuit = transpile(self.parameterized_circuit, self.backend)
        self.parametrized_circuit_depth = depth

    def sample_cost_landscape(
        self, angles={"gamma": [0, 2 * np.pi, 20], "beta": [0, 2 * np.pi, 20]}
    ):
        """
        Evaluates the cost function (landscape) at depth p = 1 for a grid of angles.

        The grid search is performed in the *vanilla (symmetric) subspace*: for every
        ``(gamma, beta)`` grid point, **all** gamma parameters for that layer are set
        equal to ``gamma`` and **all** beta parameters are set equal to ``beta``.
        This restriction keeps the search 2-D regardless of the total number of
        multi-angle parameters, while ensuring that the best grid point lies in the
        same landscape as vanilla QAOA, providing a valid warm-start for both single-
        and multi-angle ansätze.

        **Sequential mode**: If `sequential` is set to True, CVaR, variance and
        max/min cost are calculated sequentially for each (gamma, beta) combination
        and stored in `Exp_sampled_p1`, `Var_sampled_p1`, `MaxCost_sampled_p1`,
        and `MinCost_sampled_p1`.

        **Batch mode**: If `sequential` is set to False, all parameter binds are
        prepared and a single batched job is submitted to the backend.

        Args:
            angles (dict[str, list], optional): Grid specification for gamma and beta,
                each as ``[start, stop, num]``.
                Defaults to ``{"gamma": [0, 2π, 20], "beta": [0, 2π, 20]}``.

        Raises:
            NotImplementedError: If the backend is not local or does not support
                batched parameter binds.
        """
        self.landscape_p1_angles = angles
        logger = LOG.bind(func=self.sample_cost_landscape.__name__)
        logger.info("Calculating energy landscape for depth p=1...")

        depth = 1

        tmp = angles["gamma"]
        self.gamma_grid = np.linspace(tmp[0], tmp[1], tmp[2], endpoint=False)
        tmp = angles["beta"]
        self.beta_grid = np.linspace(tmp[0], tmp[1], tmp[2], endpoint=False)

        if self.backend.configuration().local:
            if self.sequential:
                expectations = []
                variances = []
                maxcosts = []
                mincosts = []

                self.createParameterizedCircuit(depth)
                logger.info("Executing sample_cost_landscape")
                for b in range(angles["beta"][2]):
                    for g in range(angles["gamma"][2]):
                        gamma = self.gamma_grid[g]
                        beta = self.beta_grid[b]
                        # Build angle array in the symmetric (vanilla) subspace: all
                        # gamma parameters equal to the grid value and all beta
                        # parameters equal to the grid value.  This ensures the grid
                        # search explores the vanilla subspace so the best grid point
                        # is a valid warm-start for both vanilla and multi-angle ansätze.
                        angle_array = (
                            [0.0] * self.n_init
                            + [gamma] * self.n_gamma
                            + [beta] * self.n_beta
                        )
                        params = self.getParametersToBind(
                            angle_array, self.parametrized_circuit_depth, asList=True
                        )

                        job = self.backend.run(
                            self.parameterized_circuit,
                            noise_model=self.noisemodel,
                            shots=self.shots,
                            parameter_binds=[params],
                            optimization_level=0,
                            memory=self.memory,
                        )

                        jres = job.result()
                        counts_list = jres.get_counts()

                        self.stat.reset()
                        for string in counts_list:
                            # qiskit binary strings use little endian encoding, but our cost function expects big endian encoding. Therefore, we reverse the order
                            cost = self.problem.cost(string[::-1])
                            self.stat.add_sample(
                                cost, counts_list[string], string[::-1]
                            )

                        expectations.append(self.stat.get_CVaR())
                        variances.append(self.stat.get_Variance())
                        maxcosts.append(self.stat.get_max())
                        mincosts.append(self.stat.get_min())

                angles = self.landscape_p1_angles
                self.Exp_sampled_p1 = -np.array(expectations).reshape(
                    angles["beta"][2], angles["gamma"][2]
                )
                self.Var_sampled_p1 = np.array(variances).reshape(
                    angles["beta"][2], angles["gamma"][2]
                )
                self.MaxCost_sampled_p1 = -np.array(maxcosts).reshape(
                    angles["beta"][2], angles["gamma"][2]
                )
                self.MinCost_sampled_p1 = -np.array(mincosts).reshape(
                    angles["beta"][2], angles["gamma"][2]
                )
                logger.info("Done measurement")
            else:
                self.createParameterizedCircuit(depth)
                gamma = [None] * angles["beta"][2] * angles["gamma"][2]
                beta = [None] * angles["beta"][2] * angles["gamma"][2]

                counter = 0
                for b in range(angles["beta"][2]):
                    for g in range(angles["gamma"][2]):
                        gamma[counter] = self.gamma_grid[g]
                        beta[counter] = self.beta_grid[b]
                        counter += 1

                parameters = {}
                # All gamma parameters get the same grid values (vanilla subspace).
                # This ensures the 2-D grid search is equivalent to vanilla QAOA,
                # giving a valid warm-start for both single- and multi-angle ansätze.
                for i in range(self.n_gamma):
                    parameters[self.gamma_params[0][i]] = gamma
                # All beta parameters get the same grid values (vanilla subspace).
                for i in range(self.n_beta):
                    parameters[self.beta_params[0][i]] = beta

                logger.info("Executing sample_cost_landscape")
                logger.info(f"parameters: {len(parameters)}")
                job = self.backend.run(
                    self.parameterized_circuit,
                    noise_model=self.noisemodel,
                    shots=self.shots,
                    parameter_binds=[parameters],
                    optimization_level=0,
                    memory=self.memory,
                )
                logger.info("Done execute")
                self.measurementStatistics(job)
                logger.info("Done measurement")

        else:
            raise NotImplementedError

        logger.info("Calculating Energy landscape done")

    def measurementStatistics(self, job):
        # """
        # implements a function for expectation value and variance

        # :param job: job instance derived from BaseJob
        # :return: expectation and variance, if the job is a list
        # """
        """
        Processes the results of a job to extract measurement statistics -- CVar, variance and max/min costs.

        Extracts the job result and retrieves the counts of measurement outcomes (stored in `last_hist` for future reference), determining whether the result is a dictionary
        (single execution) or a list of dictionaries (batched executions with parameter sweeps).
        - If the result is a list of dictionaries, it computes the measurement statistics for each set of counts, which are stored in `Exp_sampled_p1`, `Var_sampled_p1`, `MaxCost_sampled_p1`, and `MinCost_sampled_p1`.
        - If the result is a single dictionary, it computes the measurement statistics for that single set of counts, storing the results in `stat`.

        The function also handles the memory of measurement outcomes and stores them in `memory_lists` if memory is enabled
        (`memorysize` > 0).

        Args:
            job (quiskit.providers.BaseJob): Executed Quiskit job instance containing the results of the quantum circuit execution.

        """
        jres = job.result()
        counts_list = jres.get_counts()

        if self.memorysize > 0:
            for i, _ in enumerate(jres.results):
                memory_list = jres.get_memory(experiment=i)
                if self.memorysize > 0:
                    for measurement in memory_list:
                        self.memory_lists.append(
                            [measurement, self.problem.cost(measurement[::-1])]
                        )
                        self.memorysize -= 1
                        if self.memorysize < 1:
                            break

        self.last_hist = counts_list

        if isinstance(counts_list, list):
            expectations = []
            variances = []
            maxcosts = []
            mincosts = []
            for i, counts in enumerate(counts_list):
                self.stat.reset()
                for string in counts:
                    # qiskit binary strings use little endian encoding, but our cost function expects big endian encoding. Therefore, we reverse the order
                    cost = self.problem.cost(string[::-1])
                    self.stat.add_sample(cost, counts[string], string[::-1])
                expectations.append(self.stat.get_CVaR())
                variances.append(self.stat.get_Variance())
                maxcosts.append(self.stat.get_max())
                mincosts.append(self.stat.get_min())
            angles = self.landscape_p1_angles
            self.Exp_sampled_p1 = -np.array(expectations).reshape(
                angles["beta"][2], angles["gamma"][2]
            )
            self.Var_sampled_p1 = np.array(variances).reshape(
                angles["beta"][2], angles["gamma"][2]
            )
            self.MaxCost_sampled_p1 = -np.array(maxcosts).reshape(
                angles["beta"][2], angles["gamma"][2]
            )
            self.MinCost_sampled_p1 = -np.array(mincosts).reshape(
                angles["beta"][2], angles["gamma"][2]
            )
        else:
            for string in counts_list:
                # qiskit binary strings use little endian encoding, but our cost function expects big endian encoding. Therefore, we reverse the order
                cost = self.problem.cost(string[::-1])
                self.stat.add_sample(cost, counts_list[string], string[::-1])

    def optimize(
        self,
        depth,
        angles={"gamma": [0, 2 * np.pi, 20], "beta": [0, 2 * np.pi, 20]},
    ):
        """
        Runs the optimization process for the QAOA algorithm up to a specified depth.

        Drives the core iterative optimization loop of the QAOA algorithm up to a
        specified depth `p`, by incrementally building the circuit one layer at a time
        until the desired depth is reached.

        - At depth p=1 a **2-D grid search** is performed to find a good starting point.
          The grid search is always in the *vanilla (symmetric) subspace*: all gamma
          parameters for that layer are set equal to the grid gamma value and all beta
          parameters are set equal to the grid beta value.  This guarantees that the
          warm-start for multi-angle and orbit ansätze lies in the same landscape as
          vanilla QAOA.

        - At depth p>1, two strategies are available depending on ``self.interpolate``:

          * ``interpolate=True`` (default): uses the INTERP heuristic to produce a
            smooth initial guess by interpolating the optimal angles from depth p-1.
            Works well for vanilla QAOA but can give non-monotonic approximation ratios
            for multi-angle / orbit ansätze.

          * ``interpolate=False``: uses a **layer-by-layer grid search** via
            :meth:`_grid_search_layer`.  The best angles from depth p-1 are *locked*
            and a 2-D grid search is performed over the new layer's parameters in the
            vanilla subspace.  Because the grid includes (gamma=0, beta=0) — which
            adds an identity layer reproducing the depth-(p-1) result — the initial
            cost at depth p is guaranteed to be ≤ cost at depth p-1, ensuring a
            monotonically increasing approximation ratio.  Recommended for free and
            orbit ansätze.

        At each depth, the optimization results are stored in `optimization_results`.
        Measurement statistics are collected and stored in `samplecount_hists`.

        Args:
            depth (int): The maximum depth p to which the optimization should be run.
            angles (dict, optional): Dictionary specifying the grid search range for
                gamma and beta, each as ``[start, stop, num]``.
                Defaults to ``{"gamma": [0, 2π, 20], "beta": [0, 2π, 20]}``.
        """
        ## run local optimization by iteratively increasing the depth until depth p is reached
        while self.current_depth < depth:
            n_gamma = self.problem.get_num_parameters()
            n_beta = self.mixer.get_num_parameters()
            n_init = self.initialstate.get_num_parameters()
            n_per_layer = n_gamma + n_beta

            if self.current_depth == 0:
                if self.Exp_sampled_p1 is None:
                    self.sample_cost_landscape(angles=angles)
                ind_Emin = np.unravel_index(
                    np.argmin(self.Exp_sampled_p1, axis=None), self.Exp_sampled_p1.shape
                )
                # Build initial angles with init params (zeros) followed by first-layer params.
                # For multi-angle QAOA (n_gamma > 1 or n_beta > 1), the best vanilla angles
                # are broadcast to all parameters of each type as a warm-start.
                gamma_best = self.gamma_grid[ind_Emin[1]]
                beta_best = self.beta_grid[ind_Emin[0]]
                angles0 = np.array(
                    [0.0] * n_init
                    + [gamma_best] * n_gamma
                    + [beta_best] * n_beta
                )
            else:
                best_angles = self.get_angles(self.current_depth)

                if self.interpolate:
                    angles0 = self.interp(best_angles)
                else:
                    # Layer-by-layer grid search: lock previous layers at their
                    # best angles and do a 2-D grid search over the new layer.
                    # The grid includes (0, 0) which is equivalent to the depth
                    # p-1 circuit, so cost(p) ≤ cost(p-1) is guaranteed.
                    angles0 = self._grid_search_layer(best_angles, angles)

            self.optimization_results[self.current_depth + 1] = OptResult(
                self.current_depth + 1
            )
            # Create parameterized circuit at new depth
            new_depth = int((len(angles0) - n_init) / n_per_layer)
            self.createParameterizedCircuit(new_depth)

            res = self.local_opt(angles0)

            self.samplecount_hists[self.current_depth + 1] = self.last_hist

            self.optimization_results[self.current_depth + 1].compute_best_index()

            LOG.info(
                f"cost(depth { self.current_depth + 1} = {res.fun}",
                func=self.optimize.__name__,
            )

            if self.flip:
                hist = self.samplecount_hists[self.current_depth + 1]
                string = max(hist, key=hist.get)
                boosted = self.flipper.boost_samples(
                    problem=self.problem,
                    string=string,
                )
                string_xor = self.flipper.xor(string, boosted)
                self.bitflips.append(string_xor)

            self.current_depth += 1

        if self.post:
            samples = self.samplecount_hists[self.current_depth]
            post_processing(self, samples=samples, K=self.post)
            self.Exp_post_processed = -self.stat.get_CVaR()
            self.Var_post_processed = self.stat.get_Variance()

    def local_opt(self, angles0):
        """
        Performs classical local optimization of the QAOA circuit at a given depth by minimizing the loss function using the specified optimizer in `optimizer`.

        Args:
            angles0 (list): Initial guess for the angles to be optimized.

        Returns:
            res (OptimizationResult): The result of the optimization process, containing the optimized angles and the minimum loss value.

        Raises:
            TypeError: If the optimizer requires fidelity and it is not provided.
        """

        try:
            opt = self.optimizer[0](**self.optimizer[1])
        except TypeError as e:  ### QNSPSA needs fidelity
            self.isQNSPSA = True
            self.optimizer[1]["fidelity"] = self.optimizer[0].get_fidelity(
                self.parameterized_circuit, sampler=_SamplerV2()
            )
            opt = self.optimizer[0](**self.optimizer[1])
        res = opt.minimize(self.loss, x0=angles0)
        if self.isQNSPSA:
            self.optimizer[1].pop("fidelity")
        return res

    def loss(self, angles):
        """
        Loss function.

        Args:
            angles (list): List of angles (gamma and beta).

        Returns:
            float: The negative expected value of the cost function (CVaR) for the given angles.
            This is used as the objective function to be minimized during optimization.

        Raises:
            NotImplementedError: If the backend is not local or not implemented.
        """

        circuit = None
        n_target = self.shots
        self.stat.reset()
        shots_taken = 0
        shots = self.shots

        for i in range(3):  # this loop is used only used if precision is set
            if self.backend.configuration().local:
                params = self.getParametersToBind(
                    angles, self.parametrized_circuit_depth, asList=True
                )
                job = self.backend.run(
                    self.parameterized_circuit,
                    noise_model=self.noisemodel,
                    shots=shots,
                    parameter_binds=[params],
                    optimization_level=0,
                    memory=self.memory,
                )
            else:
                raise NotImplementedError
            shots_taken += shots
            self.measurementStatistics(job)
            if self.precision is None:
                break
            else:
                v = self.stat.get_Variance()
                shots = int((np.sqrt(v) / self.precision) ** 2) - shots_taken
                if shots <= 0:
                    break

        self.optimization_results[self.current_depth + 1].add_iteration(
            angles.copy(), self.stat, shots_taken
        )

        return -self.stat.get_CVaR()

    def getParametersToBind(self, angles, depth, asList=False):
        """
        Utility function to structure the parameterized parameter values
        so that they can be applied/bound to the parameterized circuit.

        The flat angle array has the format:
            [init_0, ..., init_{n_init-1},
             gamma_{0,0}, ..., gamma_{0,n_gamma-1}, beta_{0,0}, ..., beta_{0,n_beta-1},
             gamma_{1,0}, ..., gamma_{1,n_gamma-1}, beta_{1,0}, ..., beta_{1,n_beta-1},
             ...,
             gamma_{p-1,0}, ..., gamma_{p-1,n_gamma-1}, beta_{p-1,0}, ..., beta_{p-1,n_beta-1}]

        For the default single-parameter case (n_init=0, n_gamma=1, n_beta=1), this reduces
        to the classic alternating format: [gamma_0, beta_0, gamma_1, beta_1, ...].

        Args:
            angles (array-like): 1D flat array of parameter values.
            depth (int): Circuit depth.
            asList (bool): Boolean that specify if the values in the dict should be a list or not (required for batching).

        Returns:
            dict: Dictionary mapping each QAOA parameter to its corresponding value.
        """
        n_per_layer = self.n_gamma + self.n_beta
        assert len(angles) == self.n_init + depth * n_per_layer

        params = {}
        offset = 0

        # Bind initial state parameters
        for i in range(self.n_init):
            if asList:
                params[self.init_params[i]] = [angles[offset + i]]
            else:
                params[self.init_params[i]] = angles[offset + i]
        offset += self.n_init

        # Bind gamma and beta parameters layer by layer
        for d in range(depth):
            for i in range(self.n_gamma):
                if asList:
                    params[self.gamma_params[d][i]] = [angles[offset + i]]
                else:
                    params[self.gamma_params[d][i]] = angles[offset + i]
            offset += self.n_gamma

            for i in range(self.n_beta):
                if asList:
                    params[self.beta_params[d][i]] = [angles[offset + i]]
                else:
                    params[self.beta_params[d][i]] = angles[offset + i]
            offset += self.n_beta

        return params

    def interp(self, angles):
        """
        INTERP heuristic/linear interpolation for initial parameters
        when going from depth p to p+1 (https://doi.org/10.1103/PhysRevX.10.021067).

        Operates on the full flat angle array in the format:
            [init_0, ..., init_{n_init-1}, gamma_{0,0}, ..., beta_{p-1,n_beta-1}]

        Initial state parameters are kept unchanged. The INTERP heuristic is applied
        independently to each gamma and beta parameter index across layers.

        Args:
            angles (array-like): Full flat array of parameters at depth p.

        Returns:
            np.ndarray: Interpolated parameters for depth p+1 in the same flat format.
        """
        n_per_layer = self.n_gamma + self.n_beta
        depth = (len(angles) - self.n_init) // n_per_layer

        # Keep initial state parameters unchanged
        init_part = list(angles[: self.n_init])

        # Reshape layer params to [depth, n_gamma + n_beta]
        layer_angles = np.array(angles[self.n_init :]).reshape(depth, n_per_layer)

        # Apply INTERP heuristic independently to each parameter index
        result_layers = np.zeros((depth + 1, n_per_layer))
        for i in range(n_per_layer):
            param_vals = layer_angles[:, i]
            tmp = np.zeros(depth + 2)
            tmp[1:-1] = param_vals
            w = np.arange(0, depth + 1)
            result_layers[:, i] = w / depth * tmp[:-1] + (depth - w) / depth * tmp[1:]

        return np.concatenate([init_part, result_layers.flatten()])

    def _eval_cost(self, angle_array):
        """
        Evaluate the expected cost (CVaR) for a specific angle array without
        recording the result in ``optimization_results``.

        Intended for use during grid searches where many candidate points are
        evaluated cheaply without polluting the optimizer trajectory.

        Args:
            angle_array (np.ndarray): Flat parameter array whose length must
                be consistent with the current ``parametrized_circuit_depth``.

        Returns:
            float: Negative expected cost (CVaR), i.e. the value to minimise.

        Raises:
            NotImplementedError: If the backend is not local.
        """
        if not self.backend.configuration().local:
            raise NotImplementedError
        params = self.getParametersToBind(
            angle_array, self.parametrized_circuit_depth, asList=True
        )
        job = self.backend.run(
            self.parameterized_circuit,
            noise_model=self.noisemodel,
            shots=self.shots,
            parameter_binds=[params],
            optimization_level=0,
            memory=self.memory,
        )
        jres = job.result()
        counts = jres.get_counts()
        self.stat.reset()
        for string in counts:
            cost = self.problem.cost(string[::-1])
            self.stat.add_sample(cost, counts[string], string[::-1])
        return -self.stat.get_CVaR()

    def _grid_search_layer(self, prev_angles, angles):
        """
        Grid search over a single new layer's parameters with all previous
        layers locked at ``prev_angles``.

        For each ``(gamma, beta)`` point on the 2-D grid the cost is evaluated
        with layers ``1 … p-1`` fixed to ``prev_angles`` and the new layer
        ``p`` initialised in the *symmetric (vanilla) subspace*: all
        ``n_gamma`` gamma parameters are set equal to ``gamma`` and all
        ``n_beta`` beta parameters are set equal to ``beta``.

        Because the grid always includes the point ``(0, 0)`` — which adds an
        identity layer and therefore reproduces the depth-``(p-1)`` circuit —
        the best grid cost satisfies ``cost_grid ≤ cost(p-1)``.  The
        subsequent full local optimisation can only improve on this, so
        ``cost(p) ≤ cost(p-1)`` is guaranteed, giving a monotonically
        increasing approximation ratio.

        Args:
            prev_angles (np.ndarray): Flat angle array optimised at depth p-1,
                with length ``n_init + (p-1) * (n_gamma + n_beta)``.
            angles (dict): Grid specification in the form
                ``{"gamma": [lo, hi, n], "beta": [lo, hi, n]}``.

        Returns:
            np.ndarray: Initial angle array of length
                ``n_init + p * (n_gamma + n_beta)`` for the depth-p local
                optimisation.

        Raises:
            NotImplementedError: If the backend is not local.
        """
        n_per_layer = self.n_gamma + self.n_beta
        new_depth = (len(prev_angles) - self.n_init) // n_per_layer + 1
        self.createParameterizedCircuit(new_depth)

        gamma_grid = np.linspace(
            angles["gamma"][0], angles["gamma"][1], angles["gamma"][2], endpoint=False
        )
        beta_grid = np.linspace(
            angles["beta"][0], angles["beta"][1], angles["beta"][2], endpoint=False
        )

        logger = LOG.bind(func=self._grid_search_layer.__name__)
        logger.info(
            f"Layer grid search at depth {new_depth}: "
            f"{len(gamma_grid)}×{len(beta_grid)} points"
        )

        best_cost = np.inf
        best_angles = None

        for beta_val in beta_grid:
            for gamma_val in gamma_grid:
                new_layer = np.array(
                    [gamma_val] * self.n_gamma + [beta_val] * self.n_beta
                )
                candidate = np.concatenate([prev_angles, new_layer])
                cost = self._eval_cost(candidate)
                if cost < best_cost:
                    best_cost = cost
                    best_angles = candidate.copy()

        logger.info(f"Layer grid search done, best cost: {-best_cost:.6f}")
        return best_angles

    def hist(self, angles, shots):
        """
        Executes the QAOA circuit for a given set of angles and returns the corrected histogram of bitstring outcomes.

        Args:
            angles (array-like): Flat list or array of parameter values in the multi-angle format.
            shots (int): Number of circuit executions (samples) to collect.

        Returns:
            dict: Dictionary of measurement outcomes (bit-strings, little-endian format) and their counts for the given angles at the current depth.

        Raises:
            NotImplementedError: If the backend is not local or not implemented.
        """
        n_gamma = self.problem.get_num_parameters()
        n_beta = self.mixer.get_num_parameters()
        n_init = self.initialstate.get_num_parameters()
        depth = int((len(angles) - n_init) / (n_gamma + n_beta))
        self.createParameterizedCircuit(depth)

        params = self.getParametersToBind(angles, depth, asList=True)
        if self.backend.configuration().local:
            job = self.backend.run(
                self.parameterized_circuit,
                shots=shots,
                parameter_binds=[params],
                optimization_level=0,
                memory=self.memory,
            )
        else:
            raise NotImplementedError

        # Qiskit uses big endian encoding, cost function uses litle endian encoding.
        # Therefore the string is reversed before passing it to the cost function.
        hist = job.result().get_counts()
        corrected_hist = {}
        for string in hist:
            corrected_hist[string[::-1]] = hist[string]
        return corrected_hist

    def random_init(self, gamma_bounds, beta_bounds, depth):
        """
        Enforces the bounds of gamma and beta based on the graph type.

        Args:
            gamma_bounds (tuple): Parameter bound tuple (min,max) for gamma.
            beta_bounds (tuple): Parameter bound tuple (min,max) for beta.
            depth (int): The depth of the QAOA circuit.

        Returns:
            np.ndarray: On the form (gamma_1, beta_1, gamma_2, ...., gamma_d, beta_d).
        """
        gamma_list = np.random.uniform(gamma_bounds[0], gamma_bounds[1], size=depth)
        beta_list = np.random.uniform(beta_bounds[0], beta_bounds[1], size=depth)
        initial = np.empty((gamma_list.size + beta_list.size,), dtype=gamma_list.dtype)
        initial[0::2] = gamma_list
        initial[1::2] = beta_list
        return initial

    def get_optimal_solutions(self):
        """
        Iterates through all optimization result objects looking for the bit-string(s) that gave optimal cost value.

        Returns:
            list: List with the obtained best bit-strings.
        """
        best_sols = []
        best_costs = []
        for i in self.optimization_results:
            best_sols_i, best_cost_i = self.optimization_results[i].get_best_solution()
            best_sols.append(best_sols_i)
            best_costs.append(best_cost_i)

        best_iterations = np.where(best_costs == np.min(best_costs))[0]
        opt_sols = []
        for i in best_iterations:
            opt_sols.append(best_sols[i])
        opt_sols = [item for sublist in opt_sols for item in sublist]
        return np.unique(opt_sols)
