Source code for semi_cr.core.lab.station.routing_v1.quantities

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Protocol

import networkx as nx
import numpy as np
from qm import qua

if TYPE_CHECKING:
    from semi_cr.core.lab.pins.base import BasePinChannel
    from semi_cr.core.lab.pins.pins import RFInput


from qcodes.instrument_drivers.Keysight.N52xx import KeysightPNAPort
from qualang_tools.units import unit
from quam.components.pulses import SquareReadoutPulse

from semi_cr.core.lab.drivers.opx_config import OPX1000
from semi_cr.core.lab.station.routing_v1.models import RouteRequest
from semi_cr.core.lab.station.routing_v1.queries import find_route
from semi_cr.core.lab.station.routing_v1.resolution import (
    find_backend_object_for_pin,
    find_backend_parameter_for_pin,
)
from semi_cr.core.lab.utils.quam import update_rf_device_amplitude, update_rf_device_frequency

u = unit(coerce_to_integer=True)

# class RoutedElectricalQuantity:
#     quantity: Literal["voltage", "current"]

#     @property
#     def source(self): ...
#     @property
#     def meter(self): ...
#     @property
#     def readback_parameter(self): ...

# class RoutedVoltage(RoutedElectricalQuantity):
#     quantity = "voltage"

# class RoutedCurrent(RoutedElectricalQuantity):
#     quantity = "current"




class RoutedVoltage:
    def __init__(
        self,
        graph: nx.MultiDiGraph,
        pin: BasePinChannel,
        modality: str,  # "dc", "ac", "rf", "pulse", ...
    ):
        self.graph = graph
        self.pin = pin
        self.modality = modality

        self._source = None
        self._readback = None
        self._meter = None

    @property
    def source(self):

        if self._source is None:
            self._source = find_backend_parameter_for_pin(
                graph=self.graph,
                device=self.pin.parent,
                pin_name=self.pin.short_name,
                role="source",
                quantity="voltage",
                modality=self.modality,
                parameter_name="voltage",
            )

        return self._source

    @property
    def readback_parameter(self):

        if self._readback is None:
            self._readback = find_backend_parameter_for_pin(
                graph=self.graph,
                device=self.pin.parent,
                pin_name=self.pin.short_name,
                role="source",
                quantity="voltage",
                modality=self.modality,
                parameter_name="voltage_readback",
            )

        return self._readback

    @property
    def meter(self):

        if self._meter is None:
            self._meter = find_backend_parameter_for_pin(
                graph=self.graph,
                device=self.pin.parent,
                pin_name=self.pin.short_name,
                role="meter",
                quantity="voltage",
                modality=self.modality,
                parameter_name="voltage",
            )

        return self._meter



    def set(self, value: float) -> None:
        self.source.set(value)




    def get(self):
        return self.source.get()




    def readback(self) -> float:
        return float(self.readback_parameter.get())




    def measure(self) -> float:
        return float(self.meter.get())




    def invalidate(self) -> None:
        """
        Discard cached backend parameters.

        The next operation resolves the route again. This is needed after
        station reconnection, graph reconstruction, or configurable-routing
        changes.
        """
        self._source = None
        self._readback = None
        self._meter = None






class RoutedCurrent:
    def __init__(
        self,
        graph: nx.MultiDiGraph,
        pin: BasePinChannel,
        modality: str,  # "dc", "ac", "rf", "pulse", ...
    ):
        self.graph = graph
        self.pin = pin
        self.modality = modality

        self._source = None
        self._readback = None
        self._meter = None

    @property
    def source(self):
        if self._source is None:
            self._source = find_backend_parameter_for_pin(
                graph=self.graph,
                device=self.pin.parent,
                pin_name=self.pin.short_name,
                role="source",
                quantity="current",
                modality=self.modality,
                parameter_name="current",
            )

        return self._source

    @property
    def readback_parameter(self):
        if self._readback is None:
            self._readback = find_backend_parameter_for_pin(
                graph=self.graph,
                device=self.pin.parent,
                pin_name=self.pin.short_name,
                role="source",
                quantity="current",
                modality=self.modality,
                parameter_name="current_readback",
            )

        return self._readback

    @property
    def meter(self):
        if self._meter is None:
            self._meter = find_backend_parameter_for_pin(
                graph=self.graph,
                device=self.pin.parent,
                pin_name=self.pin.short_name,
                role="meter",
                quantity="current",
                modality=self.modality,
                parameter_name="current",
            )

        return self._meter



    def set(self, value: float) -> None:
        self.source.set(value)




    def get(self):
        return self.source.get()




    def readback(self) -> float:
        return float(self.readback_parameter.get())




    def measure(self) -> float:
        return float(self.meter.get())




    def invalidate(self) -> None:
        """
        Discard cached backend parameters.

        The next operation resolves the route again. This is needed after
        station reconnection, graph reconstruction, or configurable-routing
        changes.
        """
        self._source = None
        self._readback = None
        self._meter = None






@dataclass
class CurrentFromTransimpedanceVoltage:
    voltage: Any
    transimpedance_gain: Callable[[], float]



    def measure(self) -> float:
        voltage = self.voltage.measure()
        gain = float(self.transimpedance_gain())

        return voltage / gain






class RoutedRFDrive:
    def __init__(
        self,
        graph: nx.MultiDiGraph,
        pin: RFInput,
        modality: str = "rf",
    ):
        self.graph = graph
        self.pin = pin
        self.modality = modality
        self._source = None
        self.operations = {
            "readout": SquareReadoutPulse(
                length=1000,
                amplitude=0.1,
            )
        }

    @property
    def source(self):
        if self._source is None:
            self._source = find_backend_object_for_pin(
                graph=self.graph,
                device=self.pin.parent,
                pin_name=self.pin.short_name,
                role="source",
                modality=self.modality,
            )
        return self._source

    @property
    def port(self):
        return self.source.port

    @property
    def intermediate_frequency(self):
        return self.source.intermediate_frequency

    @property
    def frequency(self):
        return self.source.frequency

    @property
    def amplitude(self):
        return self.source.amplitude

    @property
    def readout_pulse(self):
        return self.operations["readout"]



    def update_frequency(self, frequency):
        self.pin.quam.channel.update_frequency(frequency)






class RoutedIQReceiver:
    def __init__(
        self,
        graph: nx.MultiDiGraph,
        pin: RFInput,
        modality: str = "rf",
    ):
        self.graph = graph
        self.pin = pin
        self.modality = modality
        self._receiver = None

    @property
    def receiver(self):
        if self._receiver is None:
            self._receiver = find_backend_object_for_pin(
                graph=self.graph,
                device=self.pin.parent,
                pin_name=self.pin.short_name,
                role="receiver",
                modality=self.modality,
            )
        return self._receiver

    @property
    def port(self):
        return self.receiver.port

    @property
    def time_of_flight(self):
        return self.receiver.time_of_flight





class RoutedRFFrequency:
    def __init__(
        self,
        graph,
        pin,
        modality="rf",
    ):
        self.graph = graph
        self.pin = pin
        self.modality = modality



    def set(self, frequency):
        update_rf_device_frequency(frequency, self.pin)






class RoutedPulseAmplitude:
    def __init__(
        self,
        graph: nx.MultiDiGraph,
        pin: BasePinChannel,
        modality="rf",
    ):
        self.graph = graph
        self.pin = pin
        self.modality = modality
        self._amplitude_scale: Any = 1.0



    def set(self, amplitude):
        update_rf_device_amplitude(amplitude, self.pin)
        # Only update local state if the delegated operation succeeded.
        self._amplitude_scale = amplitude




    def get(self):
        """
        Return the current QUA amplitude scale.

        This is protocol-construction state, not a physical hardware
        readback.
        """
        return self._amplitude_scale




    def reset(self) -> None:
        self._amplitude_scale = 1.0




# FOR KEYSIGHT VNA




class RFFrequencyController(Protocol):


    def set(

        self,
        frequency: float,
    ) -> None: ...





class RFPowerController(Protocol):


    def set(

        self,
        power: float,
    ) -> None: ...





class OPXRFFrequencyController:
    def __init__(
        self,
        rf_in,
        settling_time: float = 1e-6,
    ):
        self.rf_in = rf_in
        self.settling_time = settling_time



    def set(
        self,
        frequency: float,
    ) -> None:

        self.rf_in.update_frequency(
            frequency,
        )

        qua.wait(
            int(self.settling_time / 4e-9),
            self.rf_in.quam.id,
        )




    def configure_sweep(
        self,
        frequencies,
    ) -> None:
        raise NotImplementedError("OPX RF frequency sweeps are generated by the QUA program rather than configured as an instrument hardware sweep.")






class VNARFFrequencyController:
    def __init__(
        self,
        vna,
    ):
        self.vna = vna



    def set(
        self,
        frequency: float,
    ) -> None:

        self.vna.sweep_type("CW")
        self.vna.cw(frequency)




    def configure_sweep(
        self,
        frequencies,
    ) -> None:

        frequencies = np.asarray(
            frequencies,
            dtype=float,
        )

        if frequencies.ndim != 1:
            raise ValueError("Frequency sweep must be one-dimensional.")

        if frequencies.size < 2:
            raise ValueError("Frequency sweep requires at least two points.")

        # A LIN sweep can only represent equally spaced frequencies.
        steps = np.diff(frequencies)

        if not np.allclose(
            steps,
            steps[0],
            rtol=1e-9,
            atol=0,
        ):
            raise ValueError("VNA LIN sweep requires equally spaced frequency values.")

        self.vna.sweep_type("LIN")
        self.vna.start(frequencies[0])
        self.vna.stop(frequencies[-1])
        self.vna.points(len(frequencies))






class VNARFPowerController:
    def __init__(
        self,
        vna,
    ):
        self.vna = vna



    def set(
        self,
        power: float,
    ) -> None:

        self.vna.power(
            power,
        )






def find_route_endpoint(
    graph,
    route,
):
    for node_id in reversed(route.path):
        attrs = graph.nodes[node_id]
        obj = attrs.get("obj")

        if obj is not None:
            return obj

    raise LookupError("No runtime object found on route.")





def make_rf_frequency_controller(
    obj,
):
    if isinstance(obj, KeysightPNAPort):
        return VNARFFrequencyController(
            obj.parent,
        )

    if isinstance(obj, OPX1000):
        return OPXRFFrequencyController(
            obj,
        )

    raise TypeError(f"No RF frequency controller registered for {type(obj).__name__}.")





def find_rf_frequency_controller(
    graph,
    pin,
    modality="rf",
):
    route = find_route(
        graph,
        RouteRequest(
            pin=pin,
            role="source",
            quantity="frequency",
            modality=modality,
        ),
    )

    endpoint = find_route_endpoint(
        graph,
        route,
    )

    return make_rf_frequency_controller(
        endpoint,
    )



# Need to test this class with a measurement using OPX in Noma setup

# class RoutedRFFrequency:
#     def __init__(
#         self,
#         graph,
#         pin,
#         modality="rf",
#     ):
#         self.graph = graph
#         self.pin = pin
#         self.modality = modality

#     def controller(self):
#         return find_rf_frequency_controller(
#             graph=self.graph,
#             pin=self.pin,
#             modality=self.modality,
#         )

#     def set(
#         self,
#         frequency: float,
#     ) -> None:

#         self.controller().set(frequency)

#     def configure_sweep(
#         self,
#         frequencies,
#     ) -> None:
#         self.controller().configure_sweep(frequencies)


# ---------------------- RF Power ----------------------- #




def make_rf_power_controller(
    obj,
):
    if isinstance(obj, KeysightPNAPort):
        return VNARFPowerController(
            obj,
        )

    # if isinstance(obj, OPX1000):
    #     return OPXRFFrequencyController(
    #         obj,
    #     )

    raise TypeError(f"No RF frequency controller registered for {type(obj).__name__}.")





def find_rf_power_controller(
    graph,
    pin,
    modality="rf",
):
    route = find_route(
        graph,
        RouteRequest(
            pin=pin,
            role="source",
            quantity="power",
            modality=modality,
        ),
    )

    endpoint = find_route_endpoint(
        graph,
        route,
    )

    return make_rf_power_controller(
        endpoint,
    )





class RoutedRFPower:
    def __init__(
        self,
        graph,
        pin,
        modality="rf",
    ):
        self.graph = graph
        self.pin = pin
        self.modality = modality



    def set(
        self,
        power: float,
    ) -> None:

        controller = find_rf_power_controller(
            graph=self.graph,
            pin=self.pin,
            modality=self.modality,
        )

        controller.set(power)