<filename>src/adk/generators/instruction_generator.py
from abc import ABC, abstractmethod
from typing import Dict, Any, List, Type, Optional
import logging

from adk.type_aliases import DefaultPayloadType, InstructionType, LogEntryType, QuantumStateType
from adk.parsers.encoders import JSONEncoder


class BaseInstruction(ABC):
    template: InstructionType = {}

    def __init__(self, log_entry: LogEntryType) -> None:
        self._log_entry = log_entry

    @property
    @abstractmethod
    def payload(self) -> InstructionType:
        """Return an instruction dictionary to update an instruction's template with."""

    def transform(self) -> InstructionType:
        template = self.template.copy()
        template.update(self.payload)
        return template

    def _get_qubit(self, qubit: List[DefaultPayloadType]) -> Dict[str, DefaultPayloadType]:
        return {
            "node": qubit[0],
            "id": qubit[1]
        }

    def _get_node(self, key: str) -> Dict[str, str]:
        return {
            "node": self._get_field(key)
        }

    def _get_field(self, key: str) -> Any:
        return self._log_entry[key]

    def _get_groups(self) -> Dict[str, Dict[str, Any]]:
        groups_in = self._get_field("QGR")

        groups = {}
        if groups_in is not None:
            for group_id, properties in groups_in.items():
                state: Dict[str, QuantumStateType] = {}
                JSONEncoder.encode(state, properties["state"])
                groups[str(group_id)] = {
                    "qubits": [self._get_qubit(q) for q in properties['qubit_ids']],
                    "is_entangled": properties["is_entangled"],
                    "state": state["json"]
                }

        return groups

    @staticmethod
    @abstractmethod
    def get_instruction() -> str:
        """Return the instruction name from the simulator that should map to this instruction class."""


class ApplicationFinishedInstruction(BaseInstruction):
    template = {
        "command": "application-finished",
    }

    @property
    def payload(self) -> InstructionType:
        return {}

    @staticmethod
    def get_instruction() -> str:
        return "application_finished"


class UserMessageInstruction(BaseInstruction):
    template = {
        "command": "user-message",
        "message": None,
        "from": None
    }

    @property
    def payload(self) -> InstructionType:
        return {
            'message': self._get_field('LOG'),
            'from': self._get_node('FRM')
        }

    @staticmethod
    def get_instruction() -> str:
        return "user_msg"


class EntanglementInstruction(BaseInstruction):
    template: Dict[str, Any] = {
        "command": "entanglement",
        "action": None,
        "from": None,
        "to": None,
        "channels": None,
        "groups": None,
    }

    @property
    def payload(self) -> InstructionType:
        qubits: List[List[DefaultPayloadType]] = list(map(list, zip(self._get_field("NOD"), self._get_field("QID"))))
        return {
            'action': self._get_action(),
            'from': self._get_qubit(qubits[0]),
            'to': self._get_qubit(qubits[1]),
            'channels': self._get_field("PTH"),
            'groups': self._get_groups(),
        }

    @staticmethod
    @abstractmethod
    def _get_action() -> str:
        """EntanglementInstruction subclass has to indicate which entanglement action it represents."""

    @staticmethod
    @abstractmethod
    def get_instruction() -> str:
        """Instruction name from simulator only exists for subclasses of EntanglementInstruction."""


class EntanglementStartInstruction(EntanglementInstruction):
    @staticmethod
    def _get_action() -> str:
        return "start"

    @staticmethod
    def get_instruction() -> str:
        return "epr_EntanglementStage.START"


class EntanglementFinishInstruction(EntanglementInstruction):
    @staticmethod
    def _get_action() -> str:
        return "success"

    @staticmethod
    def get_instruction() -> str:
        return "epr_EntanglementStage.FINISH"


class ClassicalMessageInstruction(BaseInstruction):
    template = {
        "command": "classical-message",
        "message": None,
        "from": None,
        "to": None
    }

    @property
    def payload(self) -> InstructionType:
        return {
            'message': self._get_field('MSG'),
            'from': self._get_node('SEN'),
            'to': self._get_node('REC')
        }

    @staticmethod
    def get_instruction() -> str:
        return "SEND"


class ApplyGateInstruction(BaseInstruction):
    template = {
        "command": "apply-gate",
        "qubits": None,
        "gate": None,
        "parameters": None,
        "groups": None,
        "outcome": None
    }

    @property
    def payload(self) -> InstructionType:
        qubits = []
        qids = self._get_field('QID')
        if qids is not None:
            for qid in qids:
                qubits.append(
                    self._get_qubit([self._get_field("FRM"), qid])
                )
        return {
            'qubits': qubits,
            'gate': self._get_field('GAT'),
            'groups': self._get_groups(),
            'outcome': self._get_field('OUT')
        }

    @staticmethod
    def get_instruction() -> str:
        return "apply_gate"


class InstructionGenerator:
    INSTRUCTION_CLASSES = [ApplicationFinishedInstruction, ApplyGateInstruction, EntanglementStartInstruction,
                           EntanglementFinishInstruction, ClassicalMessageInstruction, UserMessageInstruction]

    INSTRUCTION_MAPPING = {}
    for instruction_class in INSTRUCTION_CLASSES:
        INSTRUCTION_MAPPING[instruction_class.get_instruction()] = instruction_class

    @staticmethod
    def generate(log_entry: LogEntryType) -> List[InstructionType]:
        instruction = log_entry['INS']
        cls: Optional[Type[BaseInstruction]] = InstructionGenerator.INSTRUCTION_MAPPING.get(instruction, None)

        instructions = []
        if cls is not None:
            instruction = cls(log_entry).transform()
            instructions.append(instruction)
        else:
            logging.debug("InstructionGenerator: unknown instruction (%s)", instruction)
        return instructions

<filename>src/adk/generators/instruction_generator.py
from abc import ABC, abstractmethod
from typing import Dict, Any, List, Type, Optional
import logging

from adk.type_aliases import DefaultPayloadType, InstructionType, LogEntryType, QuantumStateType
from adk.parsers.encoders import JSONEncoder


class BaseInstruction(ABC):
    template: InstructionType = {}

    def __init__(self, log_entry: LogEntryType) -> None:
        self._log_entry = log_entry

    @property
    @abstractmethod
    def payload(self) -> InstructionType:
        """Return an instruction dictionary to update an instruction's template with."""

    def transform(self) -> InstructionType:
        template = self.template.copy()
        template.update(self.payload)
        return template

    def _get_qubit(self, qubit: List[DefaultPayloadType]) -> Dict[str, DefaultPayloadType]:
        return {
            "node": qubit[0],
            "id": qubit[1]
        }

    def _get_node(self, key: str) -> Dict[str, str]:
        return {
            "node": self._get_field(key)
        }

    def _get_field(self, key: str) -> Any:
        return self._log_entry[key]

    def _get_groups(self) -> Dict[str, Dict[str, Any]]:
        groups_in = self._get_field("QGR")

        groups = {}
        if groups_in is not None:
            for group_id, properties in groups_in.items():
                state: Dict[str, QuantumStateType] = {}
                JSONEncoder.encode(state, properties["state"])
                groups[str(group_id)] = {
                    "qubits": [self._get_qubit(q) for q in properties['qubit_ids']],
                    "is_entangled": properties["is_entangled"],
                    "state": state["json"]
                }

        return groups

    @staticmethod
    @abstractmethod
    def get_instruction() -> str:
        """Return the instruction name from the simulator that should map to this instruction class."""


class ApplicationFinishedInstruction(BaseInstruction):
    template = {
        "command": "application-finished",
    }

    @property
    def payload(self) -> InstructionType:
        return {}

    @staticmethod
    def get_instruction() -> str:
        return "application_finished"


class UserMessageInstruction(BaseInstruction):
    template = {
        "command": "user-message",
        "message": None,
        "from": None
    }

    @property
    def payload(self) -> InstructionType:
        return {
            'message': self._get_field('LOG'),
            'from': self._get_node('FRM')
        }

    @staticmethod
    def get_instruction() -> str:
        return "user_msg"


class EntanglementInstruction(BaseInstruction):
    template: Dict[str, Any] = {
        "command": "entanglement",
        "action": None,
        "from": None,
        "to": None,
        "channels": None,
        "groups": None,
    }

    @property
    def payload(self) -> InstructionType:
        qubits: List[List[DefaultPayloadType]] = list(map(list, zip(self._get_field("NOD"), self._get_field("QID"))))
        return {
            'action': self._get_action(),
            'from': self._get_qubit(qubits[0]),
            'to': self._get_qubit(qubits[1]),
            'channels': self._get_field("PTH"),
            'groups': self._get_groups(),
        }

    @staticmethod
    @abstractmethod
    def _get_action() -> str:
        """EntanglementInstruction subclass has to indicate which entanglement action it represents."""

    @staticmethod
    @abstractmethod
    def get_instruction() -> str:
        """Instruction name from simulator only exists for subclasses of EntanglementInstruction."""


class EntanglementStartInstruction(EntanglementInstruction):
    @staticmethod
    def _get_action() -> str:
        return "start"

    @staticmethod
    def get_instruction() -> str:
        return "epr_EntanglementStage.START"


class EntanglementFinishInstruction(EntanglementInstruction):
    @staticmethod
    def _get_action() -> str:
        return "success"

    @staticmethod
    def get_instruction() -> str:
        return "epr_EntanglementStage.FINISH"


class ClassicalMessageInstruction(BaseInstruction):
    template = {
        "command": "classical-message",
        "message": None,
        "from": None,
        "to": None
    }

    @property
    def payload(self) -> InstructionType:
        return {
            'message': self._get_field('MSG'),
            'from': self._get_node('SEN'),
            'to': self._get_node('REC')
        }

    @staticmethod
    def get_instruction() -> str:
        return "SEND"


class ApplyGateInstruction(BaseInstruction):
    template = {
        "command": "apply-gate",
        "qubits": None,
        "gate": None,
        "parameters": None,
        "groups": None,
        "outcome": None
    }

    @property
    def payload(self) -> InstructionType:
        qubits = []
        qids = self._get_field('QID')
        if qids is not None:
            for qid in qids:
                qubits.append(
                    self._get_qubit([self._get_field("FRM"), qid])
                )
        return {
            'qubits': qubits,
            'gate': self._get_field('GAT'),
            'groups': self._get_groups(),
            'outcome': self._get_field('OUT')
        }

    @staticmethod
    def get_instruction() -> str:
        return "apply_gate"


class InstructionGenerator:
    INSTRUCTION_CLASSES = [ApplicationFinishedInstruction, ApplyGateInstruction, EntanglementStartInstruction,
                           EntanglementFinishInstruction, ClassicalMessageInstruction, UserMessageInstruction]

    INSTRUCTION_MAPPING = {}
    for instruction_class in INSTRUCTION_CLASSES:
        INSTRUCTION_MAPPING[instruction_class.get_instruction()] = instruction_class

    @staticmethod
    def generate(log_entry: LogEntryType) -> List[InstructionType]:
        instruction = log_entry['INS']
        cls: Optional[Type[BaseInstruction]] = InstructionGenerator.INSTRUCTION_MAPPING.get(instruction, None)

        instructions = []
        if cls is not None:
            instruction = cls(log_entry).transform()
            instructions.append(instruction)
        else:
            logging.debug("InstructionGenerator: unknown instruction (%s)", instruction)
        return instructions

Return an instruction dictionary to update an instruction's template with. Return the instruction name from the simulator that should map to this instruction class. EntanglementInstruction subclass has to indicate which entanglement action it represents. Instruction name from simulator only exists for subclasses of EntanglementInstruction.