Source code for prsctrl.measurement.settings

"""
Measurement settings for photoreflectance implemented as classes.
This new structure is "more strongly typed" than the old dictionary approach, and makes for less magic values (settings keys in the code).
It is also easier to apply the settings using the `apply(device)` method.
Every setting will also always be present and have a default value when not set explcitly.

Loading and saving from a file is also easier.

You can use the `dataclasses.asdict(<settings class>)` to convert a settings object to a dictionary.
"""
from devctrl import devices as d
from devctrl.data.spectrum.functions import nm_to_eV, eV_to_nm
from devctrl.utility.settings import SettingsClass

from dataclasses import dataclass, field
from typing import Any, Optional, Callable, Literal

from .offset import add_offset_measurements_in_time_intervals

import numpy as np

import logging
log = logging.getLogger(__name__)

_decials3 = dict(decimals=3)



def convert_range_to_list(wl_range: tuple | list | np.ndarray) -> list:
    """
    :param wl_range:
        if tuple, return list(range(*wl_range))
        if list or np.ndarray, return copy of wl_range
    """
    if isinstance(wl_range, tuple):
        wavelengths = np.arange(*wl_range).tolist()
    elif isinstance(wl_range, list):
        wavelengths = wl_range.copy()
    elif isinstance(wl_range, np.ndarray):
        wavelengths = wl_range.tolist()
    else:
        raise ValueError(f"Invalid type for : {type(wl_range)}")
    return wavelengths





@dataclass
class MeasurementParameters(SettingsClass):
    measurement_time_s: float = field(
        default=30,
        metadata=dict(
            doc="Time in seconds where data is recorded at each wavelength.",
            valid=(0, 9999999),
        ),
    )
    wait_time_s: Optional[float] = field(
        default=None,
        metadata=dict(
            doc="Time in seconds to wait for the signal to stabilize after changing the wavelength.\nSet to `0` for no wait time, and to `None` to automatically choose a value based on the lock-in time constant.",
            valid=(0, 999999),
        ),
    )
    additional_wait_time_s: float = field(
        default=0,
        metadata=dict(
            doc="Time in seconds to wait for the signal to stabilize after changing the wavelength.\nSet to `0` for no wait time, and to `None` to automatically choose a value based on the lock-in time constant.",
            valid=(0, 999999),
        ),
    )
    pump_source: str = field(
        default="none",
        metadata=dict(
            doc="String describing which pump source is used.\nIf 'led', then the values from :py:attr:`PrsSettings.led` are applied. If 'laser', then the values from :py:attr:`PrsSettings.laser` are applied.",
            valid=["laser", "led", "laser-manual", "led-manual", "none", "manual"],
        ),
    )
    pump_power_mWcm2: Optional[float] = field(
        default=None,
        metadata=dict(
            doc="Power in mW/cm^2. If ``pump_source`` is set and a conversion table exists for the source then the ``power_mW`` setting for lasers or the `current_A` setting for LEDs can be overwritten with the value corresponding to the ``power_mWcm2`` setting.",
        ),
    )
    auto_gain: bool = field(
        default=False,
        metadata=dict(
            doc="Whether to automatically adjust the gain during the measurement to stay in the optimum DC region. Not current implemented yet.",
        ),
    )
    auto_sensitivity: bool = field(
        default=True,
        metadata=dict(
            doc="Whether to automatically increase the lock-in sensitivity on overloads.",
        ),
    )
    wavelengths_nm: Optional[tuple[float,float,float] | list[float]] = field(
        default=None,
        metadata=dict(
            doc="Wavelength range in nm. Must be either a list of values or a tuple (start, stop, step). Only one of ``wavelengths_nm`` or ``energies_eV`` may be set.",
            metadatas=dict(
                tuple=dict(
                    metadatas=(_decials3|dict(default=420), _decials3|dict(default=800), _decials3|dict(default=1)),
                    labels = ("Start", "Stop", "Step"),
                ),
                list = dict(metadatas=_decials3)
            )
        ),
    )
    energies_eV: Optional[tuple[float,float,float] | list[float]] = field(
        default=None,
        metadata=dict(
            doc="Energy range in eV. Must be either a list of values or a tuple (start, stop, step). Only one of ``wavelengths_nm`` or ``energies_eV`` may be set.",
            metadatas=dict(
                tuple=dict(
                    metadatas=(_decials3|dict(default=1.55), _decials3|dict(default=2.8), _decials3|dict(default=0.005)),
                    labels=("Start", "Stop", "Step"),
                ),
                list=dict(metadatas=_decials3)
            )
        ),
    )
    offset_strategy: Literal["none", "start-end", "time-interval"] = field(
        default="time-interval",
        metadata=dict(
            doc="Which offset reduction strategy to use. Determines where offset measurements are added.",
        ),
    )
    offset_intervals_minutes: list[int] = field(
        default_factory=lambda: [5, 10, 15, 30, 60],
        metadata=dict(
            doc="How many minutes to wait after the last offset measurement. The last interval is repeated indefinetely.\nOnly applies when offset_strategy='time-interval'.\nStart and end offsets are always recorded, unless offset_strategy='none'.",
        ),
    )



    def get_wavelengths_nm(self, round_to_digits: Optional[int]=2, update_field: bool=False) -> list[float]:
        """Generate a list of wavelengths according to either the 'wavelengths_nm' or 'energies_eV' field.

        Either field may already be a list of values, or a tuple (start, stop, step size).

        :param round_to_digits: Optional number of digits to round the values to.
        :param update_field: If True, overwrite the value of 'wavelengths_nm' with the returned array and 'energies_eV' with None.
        :return: list of wavelengths in nm
        """
        if self.wavelengths_nm is not None and self.energies_eV is not None:
            raise KeyError(f"Only one of 'wavelengths_nm' or 'energies_eV' must be set, but both are.")
        if self.wavelengths_nm is not None:
            wavelengths = convert_range_to_list(self.wavelengths_nm)
        elif self.energies_eV is not None:
            wavelengths = [(v if type(v) == str else eV_to_nm(v)) for v in convert_range_to_list(self.energies_eV)]
        else:
            raise KeyError(f"One of 'wavelengths_nm' or 'energies_eV' must be set, but both are None")
        if round_to_digits is not None:
            wavelengths = [(v if type(v) == str else round(v, round_to_digits)) for v in wavelengths]

        if update_field:
            self.wavelengths_nm = wavelengths.copy()
            self.energies_eV = None
        return wavelengths




    def get_wavelengths_nm_with_offsets(self, time_per_wavelength: int, round_to_digits: Optional[int]=2, update_field: bool=False) -> list[float|str]:
        """Generate a list of wavelengths according to either the 'wavelengths_nm' or 'energies_eV' field and offset measurements.

        :param time_per_wavelength: Estimated runtime per wavelength measurements. Required when offset_strategy=="time-interval".
        :param round_to_digits: Optional number of digits to round the values to.
        :param update_field: If True, overwrite the value of 'wavelengths_nm' with the returned array and 'energies_eV' with None.
        :return: list of wavelengths in nm and offset measurements
        """
        wavelengths = self.get_wavelengths_nm(round_to_digits, update_field)
        if self.offset_strategy == "none":
            return wavelengths
        elif self.offset_strategy == "start-end":
            return ["offset_begin"] + wavelengths + ["offset_end"]
        elif self.offset_strategy == "time-interval":
            wls = add_offset_measurements_in_time_intervals(wavelengths, time_per_wavelength, self.offset_intervals_minutes)
            return wls




    def get_metadata(self) -> dict[str, Any]:
        """Return the measurements settings converted to a metadata dictionary that contains only relevant information.

        Leaves out wavelengths and energies, as those will be contained in the data as columns anyways.
        Leaves out optional values that are set to ``None``.
        """
        md = self.as_dict()
        del md["energies_eV"]
        del md["wavelengths_nm"]
        for k in list(md.keys()):
            if md[k] is None:
                del md[k]
        if self.offset_strategy != "time-interval":
            del md["offset_intervals_minutes"]
        else:
            md["offset_intervals_minutes"] = ', '.join([str(m) for m in self.offset_intervals_minutes])
        return md





    




@dataclass
class AmplifierSettings(SettingsClass):
    coupling: str = "DC"
    gain: float = 5
    mode: Optional[str] = None



    def apply(self, device: d.Amplifier):
        device.set_coupling(self.coupling)
        if self.mode is not None:
            if not hasattr(device, "set_mode"):
                raise RuntimeError(f"Can not apply amplifier mode, since amplifier of type {device.get_device_name()} has no attribute 'set_mode'")
            device.set_mode(self.mode)
        device.set_gain(self.gain)





@dataclass
class LaserSettings(SettingsClass):
    power_mW: float = 1.0
    mode: str = "TTL"



    def apply(self, device: d.Laser):
        device.set_mode(self.mode)
        device.set_power_mW(self.power_mW)






@dataclass
class LedSettings(SettingsClass):
    mode: str = "TTL"
    current_A: float = 0.1



    def apply(self, device: d.LedController):
        device.set_mode(self.mode)
        device.set_current_A(self.current_A)






@dataclass
class LockInSettings(SettingsClass):
    time_constant_s: float = field(
        default=1,
        metadata={
            "doc": "The Lock-In time constant in seconds.",
        }
    )
    sensitivity_volt: float = field(
        default=50e-6,
        metadata={
            "doc": "Lock-In sensitivity in volts.",
        }
    )
    
    filter_slope: float = field(
        default=12,
        metadata={
            "doc": "The slope of the lock-in's filter in dB/Oct.",
        }
    )
    sync_filter: float = field(
        default=1,
        metadata={
            "doc": "Whether the sync filter should be enabled (1) or disabled (0).",
            "valid": [0, 1],
        }
    )
    reserve: str = field(
        default="Normal",
        metadata={
            "doc": "Lock-in's dynamic reserve setting.",
        }
    )
    reference: str = field(
        default="External",
        metadata={
            "doc": "The source of the modulation signal.",
            "valid": ["Internal", "External"],
        }
    )

    reference_trigger: str = field(
        default="Rising Edge",
        metadata={
            "doc": "On which signal part to trigger. Influences the absolute phase value of the signal.",
            "valid": ["Rising Edge", "Falling Edge", "Sine"] ,
        }
    )
    frequency_Hz: Optional[float] = field(
        default=None,
        metadata={
            "doc": "Frequency of the internal oscillator. Only used when reference='Internal'.",
        }
    )
    sample_rate_Hz: float|str = field(
        default=128.,
        metadata={
            "doc": "The rate at samples are written to the buffer. If 'max', the maximum is determined based on the measurement duration and the buffer size. ",
        }
    )
    aux_DC: str = field(
        default="Aux In 1",
        metadata={
            "doc": "Auxiliary port name at which the DC part of the signal is measured.",
            "valid": ["Aux In 1", "Aux In 2", "Aux In 3", "Aux In 4"],
        }
    )



    def apply(self, device: d.LockInAmp):
        """Set 

        - time constant
        - sensitivity
        - filter slope
        - sync filter
        - reserve
        - reference source
        - reference trigger setting
        - frequency
        """
        device.set_time_constant_s(self.time_constant_s)
        device.set_sensitivity_volt(self.sensitivity_volt)
        device.set_filter_slope(self.filter_slope)
        device.set_sync_filter(self.sync_filter)
        if hasattr(device, "set_reserve"):
            device.set_reserve(self.reserve)
        device.set_reference(self.reference)
        device.set_reference_trigger(self.reference_trigger)
        if self.frequency_Hz is not None:
            device.set_frequency_Hz(self.frequency_Hz)






@dataclass
class MonochromatorSettings(SettingsClass):
    """
    Monochromator settings, to apply a bandwidth/resolution
    """
    resolution_nm: Optional[float] = field(
        default=1.0,
        metadata={
            "doc": "Resolution in nm.",
            "metadatas": {
                "float": {"decimals": 3, "valid": (0.001, 100, 0.1), "default": 1.0},
            }
        }
    )
    resolution_eV: Optional[float] = field(
        default=None,
        metadata={
            "doc": "Resolution in eV.",
            "metadatas": {
                "float": {"decimals": 3, "valid": (0.001, 1, 0.001), "default": 0.010},
            }
        }
    )



    def apply(self, device: d.Monochromator):
        if self.resolution_eV is not None and self.resolution_nm is not None:
            raise ValueError(f"Only one of 'resolution_eV' or 'resolution_nm' is allowed.")
        elif self.resolution_eV is None and self.resolution_nm is None:
            raise ValueError(f"One of 'resolution_eV' or 'resolution_nm' must not be 'None', but both are.")
        if self.resolution_eV is not None:
            device.set_resolution_eV(self.resolution_eV)
        else:
            device.set_resolution_nm(self.resolution_nm)





@dataclass
class ModeSettings(SettingsClass):
    pre_amplifier: Optional[AmplifierSettings] = None
    lock_in: LockInSettings = field(default_factory=LockInSettings)



    def get_metadata(self) -> dict[str, Any]:
        md = {}
        if self.pre_amplifier:
            md["pre-amplifier"] = self.pre_amplifier.as_dict()
        if self.lock_in:
            md["lock-in-amplifier"] = self.lock_in.as_dict()
        return md





@dataclass
class PrsSettings(SettingsClass):
    """All settings for a photoreflectance-type measurement.
    """
    modes: dict[str, ModeSettings] = field(
        default_factory=dict,
        metadata = dict(
            doc = """Mode key (ref/tra): mode settings""",
        )
    )

    general: MeasurementParameters = field(default_factory=MeasurementParameters)

    laser: Optional[LaserSettings] = None
    led: Optional[LedSettings] = None
    monochromators: list[MonochromatorSettings] = field(
        default_factory=list,
        metadata = dict(
            doc = "A list of monochromator settings is supported to allow for dual monochromator setups.",
        )
    )

    auto_set_init_gain: bool = field(  # TODO set to bool|int when supported in GUI
        default = False,
        metadata=dict(
            doc="Whether to automatically set the pre-amplifier gain before the measurement. If set to an integer, `True` is implied and the value is the number of sample points in the spectral range to take into account for determining the gain.",
        )
    )



    def prepare(self,
            power_density_to_laser_power_mW: Optional[Callable[[float], float]]=None,
            power_density_to_led_current_A: Optional[Callable[[float], float]]=None,
        ):
        """Apply pump power density setting.

        If the pump source is ``led`` or ``laser`` and `.MeasurementParameters.pump_power_mWcm2` is set,
        the power density is converted to the laser power in mW or led current in A using the conversion functions.
        These can be the methods of :py:class:`~prsctrl.measurement.pump_power_conversion.PowerConversion` object.

        :param power_density_to_laser_power_mW: Function to convert pump power density to laser power density.
        :param power_density_to_led_current_A: Function to convert pump power density to led current density.
        :raises: RuntimeError if
            - pump source is set to ```led`` or ``laser`` but not settings or power density are set
            - a power density is set but no power density conversion function was passed to this method

        """
        if self.general.pump_source == "laser":
            if self.laser is None and self.general.pump_power_mWcm2 is None:
                    raise RuntimeError(f"Pump source is 'laser', but no laser settings and no pump power density are set.")
            if self.general.pump_power_mWcm2 is not None:
                if power_density_to_laser_power_mW is None:
                    raise RuntimeError(f"No pump power conversion function given to convert target power density {self.general.pump_power_mWcm2} mW/cm^2 to laser power in mW.")
                power_mW = power_density_to_laser_power_mW(self.general.pump_power_mWcm2)
                log.info(f"Evaluated laser power for power density {power_density_to_led_current_A} mW/cm^2 as {power_mW} mW.")
                if self.laser is not None:
                    self.laser.power_mW = power_mW
                else:
                    self.laser = LaserSettings(power_mW=power_mW)
        if self.general.pump_source == "led":
            if self.led is None and self.general.pump_power_mWcm2 is None:
                raise RuntimeError(f"Pump source is 'led', but no led settings and no pump power density are set.")
            if self.general.pump_power_mWcm2 is not None:
                if power_density_to_led_current_A is None:
                    raise RuntimeError(f"No pump power conversion function given to convert target power density {self.general.pump_power_mWcm2} mW/cm^2 to led current in A.")
                current_A = power_density_to_led_current_A(self.general.pump_power_mWcm2)
                log.info(f"Evaluated LED current for power density {power_density_to_led_current_A} mW/cm^2 as {current_A} A.")
                if self.led is not None:
                    self.led.current_A = current_A
                else:
                    self.led = LedSettings(current_A=current_A)




    def get_metadata(self, mode: Optional[str]) -> dict[str, Any]:
        """Return the measurements settings converted to a metadata dictionary that contains only relevant information.
        """
        md = self.general.get_metadata()
        if self.led: md["led"] = self.led.as_dict()
        if self.laser: md["laser"] = self.laser.as_dict()
        if self.monochromators: md["monochromator"] = [m.as_dict() for m in self.monochromators]
        if mode in self.modes:
            md |= self.modes[mode].get_metadata()

        return md



if __name__ == "__main__":
    # Demonstrate how settings can be merged
    settings = PrsSettings(modes={"ref": ModeSettings(pre_amplfier=AmplifierSettings())})
    print(settings)
    loaded_from_file = {"general": {"measurement_time_s": 42}, "modes": {"tra": {"lock_in": {"time_constant_s": 5}}}}
    settings |= loaded_from_file
    print(settings)