Source code for ska_ost_senscalc.low.service

"""
The service layer is responsible for turning validated inputs into the relevant calculation inputs,
calling any calculation functions and collating the results.
"""
from typing import TypedDict

from ska_ost_senscalc.low.bright_source_lookup import BrightSourceCatalog
from ska_ost_senscalc.low.calculator import calculate_sensitivity
from ska_ost_senscalc.low.model import CalculatorInput, LowSpectralMode, WeightingQuery
from ska_ost_senscalc.mid_utilities import Beam, CalculatorMode
from ska_ost_senscalc.subarray import SubarrayStorage
from ska_ost_senscalc.utilities import Telescope

subarray_storage = SubarrayStorage(Telescope.LOW)

FLUX_DENSITY_THRESHOLD_JY = 10.0


class BeamSizeResponse(TypedDict):
    """
    BeamSizeResponse is a typed dictionary constrained to match the schema of
    the JSON object outlining the synthesized beam size, as contained in the
    parent JSON result of a weighting endpoint query.
    """

    beam_maj_scaled: float
    beam_min_scaled: float
    beam_pa: float


class WeightingResponse(TypedDict):
    """
    SingleWeightingResponse is a typed dictionary constrained to match the
    schema of a single weighting calculation, as performed for the main
    continuum or zoom weighting calculation and for each chunk frequency.
    Child SingleWeightingResponse JSON object for each of these calculations
    are contained in the parent JSON result
    """

    weighting_factor: float | None
    sbs_conv_factor: list[float] | float
    beam_size: list[BeamSizeResponse] | BeamSizeResponse


class SensitivityResponse(TypedDict):
    """
    SensitivityResponse is a typed dictionary constrained to match the
    schema of a single sensitivity calculation.
    """

    sensitivity: float
    units: str
    warn_msg: str | None


def convert_continuum_input_and_calculate(user_input: dict) -> SensitivityResponse:
    """
    :param user_input: A kwarg dict of the HTTP parameters sent by the user
    :return: a SensitivityResponse containing the calculated sensitivity and its units
    """
    calculator_input = CalculatorInput(
        freq_centre=user_input["freq_centre"],
        bandwidth=user_input["bandwidth_mhz"],
        num_stations=user_input["num_stations"],
        pointing_centre=user_input["pointing_centre"],
        duration=user_input["duration"],
    )

    return _get_calculation_value(calculator_input)


def convert_zoom_input_and_calculate(user_input: dict) -> SensitivityResponse:
    """
    :param user_input: A kwarg dict of the HTTP parameters sent by the user
    :return: a dict containing the calculated sensitivity and its units
    """
    calculator_input = CalculatorInput(
        freq_centre=user_input["freq_centre"],
        bandwidth=user_input["spectral_resolution_hz"] * 1e-6,  # Convert to MHz
        num_stations=user_input["num_stations"],
        pointing_centre=user_input["pointing_centre"],
        duration=user_input["duration"],
    )

    return _get_calculation_value(calculator_input)


def convert_beam_input_and_get_weighting(
    params: WeightingQuery,
) -> WeightingResponse:
    """
    Calculate the beam weighting from the given beam parameters.

    Input arguments and external state are not modified by this function.
    """
    # At PI19, MID supports different modes from LOW and names them
    # differently, hence we need conversion
    low_to_mid_mode_map = {
        LowSpectralMode.CONTINUUM: CalculatorMode.CONTINUUM,
        LowSpectralMode.LINE: CalculatorMode.LINE,
    }
    calculator_mode = low_to_mid_mode_map[params.spectral_mode]

    # only declination is required for a weighting calculation
    dec_degrees = params.pointing_centre.dec.degree

    # for zoom mode, zoom frequencies need to be set too as this
    # is operated on by the shared weighting calculator code
    if params.spectral_mode == LowSpectralMode.LINE:
        zoom_frequencies = [params.freq_centre]
    else:
        zoom_frequencies = None

    beam = Beam(
        frequency=params.freq_centre,
        zoom_frequencies=zoom_frequencies,
        dec=dec_degrees,
        weighting=params.weighting_mode,
        robustness=params.robustness,
        array_configuration=params.subarray_configuration,
        calculator_mode=calculator_mode,
        telescope=Telescope.LOW,
    )

    beam_size_response = [
        {
            "beam_maj_scaled": beam_size.beam_maj.value,
            "beam_min_scaled": beam_size.beam_min.value,
            "beam_pa": beam_size.beam_pa.value,
        }
        for beam_size in beam.beam_size()
    ]

    return {
        "weighting_factor": beam.weighting_factor(),
        "sbs_conv_factor": beam.surface_brightness_conversion_factor().value,
        "beam_size": beam_size_response,
    }


def _get_calculation_value(calculator_input: CalculatorInput) -> dict:
    result = calculate_sensitivity(calculator_input)
    response = {"sensitivity": result.sensitivity, "units": result.units}
    mwa_cat = BrightSourceCatalog(threshold_jy=FLUX_DENSITY_THRESHOLD_JY)
    if mwa_cat.check_for_bright_sources(
        calculator_input.pointing_centre, calculator_input.freq_centre
    ):
        response["warn_msg"] = (
            "The specified pointing contains at least one source brighter "
            + f"than {FLUX_DENSITY_THRESHOLD_JY} Jy. Your observation may be "
            + "dynamic range limited."
        )
    return response


def get_subarray_response():
    """
    return the appropriate subarray objects
    """
    return [
        {
            "name": subarray.name,
            "label": subarray.label,
            "n_stations": subarray.n_stations,
        }
        for subarray in subarray_storage.list()
    ]