from qblox_scheduler import SpinEdge
from qblox_scheduler.structure.model import Numbers, Parameter, SchedulerSubmodule


class OperatingPoints(SchedulerSubmodule):
    """
    Submodule defining the operating voltages for a double quantum dot system.

    These parameters are used to control the electrostatic environment
    during different operation modes (e.g. initialization, control, readout).
    """

    target_gate_voltage: float = Parameter(
        initial_value=0,
        unit="v",
        vals=Numbers(min_value=-1, max_value=1),
    )
    """Voltage applied to the parent dot gate (relative to the edge connection)."""

    ancilla_gate_voltage: float = Parameter(
        initial_value=0,
        unit="v",
        vals=Numbers(min_value=-1, max_value=1),
    )
    """Voltage applied to the child dot gate (relative to the edge connection)."""

    barrier_gate_voltage: float = Parameter(
        initial_value=0,
        unit="v",
        vals=Numbers(min_value=-1, max_value=1),
    )
    """Voltage applied to the inter-dot barrier gate, controlling tunnel coupling."""


class RampingTimes(SchedulerSubmodule):
    """Ramping times between different operating points."""

    init_to_readout: float = Parameter(
        initial_value=100e-9,
        unit="s",
        vals=Numbers(min_value=0),
    )
    """Ramp time from initialization to readout."""

    readout_to_control: float = Parameter(
        initial_value=100e-9,
        unit="s",
        vals=Numbers(min_value=0),
    )
    """Ramp time from readout to control."""

    control_to_readout: float = Parameter(
        initial_value=100e-9,
        unit="s",
        vals=Numbers(min_value=0),
    )
    """Ramp time from control to readout."""

    control_to_init: float = Parameter(
        initial_value=1e-6,
        unit="s",
        vals=Numbers(min_value=0),
    )
    """Ramp time from control to readout."""


class PsbEdge(SpinEdge):
    """Custom edge for spin qubits with Pauli spin blockade (PSB) configuration."""

    init: OperatingPoints
    control: OperatingPoints
    readout: OperatingPoints
    ramps: RampingTimes

    @property
    def port_dict(self) -> dict:
        """Generate the dict to map instrument output to device gates."""
        out_dict = {
            "target_gate_voltage": self.parent_element.ports.gate,  # type: ignore[reportAttributeAccessIssue]
            "ancilla_gate_voltage": self.child_element.ports.gate,  # type: ignore[reportAttributeAccessIssue]
            "barrier_gate_voltage": self.ports.gate,
        }
        return out_dict
