# This code is part of Qiskit.
#
# (C) Copyright IBM 2018, 2021.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.
"""
Simulator instruction to save custom internal data to results.
"""

import copy

from qiskit.circuit import Instruction
from qiskit.extensions.exceptions import ExtensionError


class SaveData(Instruction):
    """Pragma Instruction to save simulator data."""

    _directive = True
    _allowed_subtypes = set(
        ["single", "c_single", "list", "c_list", "average", "c_average", "accum", "c_accum"]
    )

    def __init__(self, name, num_qubits, label, subtype="single", params=None):
        """Create new save data instruction.

        Args:
            name (str): the name of hte save instruction.
            num_qubits (int): the number of qubits for the snapshot type.
            label (str): the key for retrieving saved data from results.
            subtype (str): the data subtype for the instruction [Default: 'single'].
            params (list or None): Optional, the parameters for instruction
                                   [Default: None].

        Raises:
            ExtensionError: if the subtype string is invalid.

        Additional Information:
            The supported subtypes are 'single', 'list', 'c_list', 'average',
            'c_average', 'accum', 'c_accum'.
        """
        if subtype not in self._allowed_subtypes:
            raise ExtensionError("Invalid data subtype for SaveData instruction.")

        if not isinstance(label, str):
            raise ExtensionError(
                f"Invalid label for save data instruction, {label} must be a string."
            )

        if params is None:
            params = {}

        super().__init__(name, num_qubits, 0, params)

        self._label = label
        self._subtype = subtype

    def assemble(self):
        """Return the QasmQobjInstruction for the intructions."""
        instr = super().assemble()
        # Use same fields as Snapshot instruction
        # so we dont need to modify QasmQobjInstruction
        instr.snapshot_type = self._subtype
        instr.label = self._label
        return instr

    def inverse(self):
        """Special case. Return self."""
        return copy.copy(self)


class SaveAverageData(SaveData):
    """Save averageble data"""

    def __init__(
        self,
        name,
        num_qubits,
        label,
        unnormalized=False,
        pershot=False,
        conditional=False,
        params=None,
    ):
        """Create new save data instruction.

        Args:
            name (str): the name of hte save instruction.
            num_qubits (int): the number of qubits for the snapshot type.
            label (str): the key for retrieving saved data from results.
            unnormalized (bool): If True return save the unnormalized accumulated
                                 or conditional accumulated data over all shot.
                                 [Default: False].
            pershot (bool): if True save a list of data for each shot of the
                            simulation rather than the average over  all shots
                            [Default: False].
            conditional (bool): if True save the average or pershot data
                                conditional on the current classical register
                                values [Default: False].
            params (list or None): Optional, the parameters for instruction
                                   [Default: None].
        """
        if pershot:
            subtype = "list"
        elif unnormalized:
            subtype = "accum"
        else:
            subtype = "average"
        if conditional:
            subtype = "c_" + subtype
        super().__init__(name, num_qubits, label, subtype=subtype, params=params)


class SaveSingleData(SaveData):
    """Save non-averagable single data type."""

    def __init__(self, name, num_qubits, label, pershot=False, conditional=False, params=None):
        """Create new save data instruction.

        Args:
            name (str): the name of the save instruction.
            num_qubits (int): the number of qubits for the snapshot type.
            label (str): the key for retrieving saved data from results.
            pershot (bool): if True save a list of data for each shot of the
                            simulation [Default: False].
            conditional (bool): if True save data conditional on the
                                current classical register values
                                [Default: False].
            params (list or None): Optional, the parameters for instruction
                                   [Default: None].
        """
        subtype = "list" if pershot else "single"
        if conditional:
            subtype = "c_" + subtype
        super().__init__(name, num_qubits, label, subtype=subtype, params=params)
