from typing import Union
from typing import List


def final_array_to_use(
    original_array, subarray_selection: Union[str, List[int]], subarrays=None
):
    """Infer IDs of telescopes and cameras with equivalent focal lengths.

    This is an helper function for utils.prod3b_array.

    Parameters
    ----------
    subarrays : `dict` [`str`, `list` [`int`]], optional
        Dictionary of subarray names linked to lists of tel_ids
        automatically extracted by telescope type.
        If set, it will extract tel_ids from there, otherwise from the custom
        list given by 'array'.
    original_array : `ctapipe.instrument.SubarrayDescription`
        Full simulated array from the first event.
    subarray_selection : {`str`, `list` [`int`]}
        Custom list of telescope IDs that the user wants to use or name of
        specific subarray.

    Returns
    -------
    tel_ids : `list` [`int`]
        List of telescope IDs to use in a format readable by ctapipe.
    cams_and_foclens : `dict` [`str`, `float`]
        Dictionary containing the IDs of the involved cameras as inferred from
        the involved telescopes IDs, together with the equivalent focal lengths
        of the telescopes.
        The camera IDs will feed both the estimators and the image cleaning.
        The equivalent focal lengths will be used to calculate the radius of
        the camera on which cut for truncated images.
    subarray : ctapipe.instrument.SubarrayDescription
        Complete subarray information of the final array/subarray selected.

    """
    if subarrays:
        tel_ids = subarrays[subarray_selection]
        selected_subarray = original_array.select_subarray(
            tel_ids, name="selected_subarray"
        )
    else:
        selected_subarray = original_array.select_subarray(
            subarray_selection, name="selected_subarray"
        )
        tel_ids = selected_subarray.tel_ids
    tel_types = selected_subarray.telescope_types
    cams_and_foclens = {
        tel_types[i]
        .camera.camera_name: tel_types[i]
        .optics.equivalent_focal_length.value
        for i in range(len(tel_types))
    }
    return set(tel_ids), cams_and_foclens, selected_subarray