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.common.beam import calculate_weighting
from ska_ost_senscalc.common.model import WeightingInput
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
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 subband frequency.
    Child SingleWeightingResponse JSON object for each of these calculations
    are contained in the parent JSON result
    """

    weighting_factor: float
    sbs_conv_factor: list[float]
    beam_size: list[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_input_and_get_weighting(params: WeightingInput) -> WeightingResponse:
    """
    Calculate the beam weighting from the given beam parameters.

    Input arguments and external state are not modified by this function.

    """
    weighting_result = calculate_weighting(params)

    return {
        "weighting_factor": weighting_result.weighting_factor,
        "sbs_conv_factor": [
            weighting_result.surface_brightness_conversion_factor[0].value
        ],
        "beam_size": [
            {
                "beam_maj_scaled": weighting_result.beam_size[0].beam_maj.value,
                "beam_min_scaled": weighting_result.beam_size[0].beam_min.value,
                "beam_pa": weighting_result.beam_size[0].beam_pa.value,
            }
        ],
    }


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()
    ]