PredictNDFilters

class PredictNDFilters(KPFFunction):
    '''Predict which ND filters should be used for simultaneous calibrations.

    Args:
        Gmag (float): The Gaia G magnitude of the target.
        Teff (float): The effective temperature of the target.
        ExpTime (float): The exposure time.
        set (bool): Set these values or just calculate?

    Functions Called:

     - `kpf.calbench.SetND`
    '''
    @classmethod
    def pre_condition(cls, args):
        check_input(args, 'Gmag', allowed_types=[int, float])
        check_input(args, 'Teff', allowed_types=[int, float])
        check_input(args, 'ExpTime', allowed_types=[int, float])

    @classmethod
    def perform(cls, args):
        tick = datetime.datetime.now()
        gmag = args.get('Gmag')
        teff = args.get('Teff')
        log.debug('Estimating Vmag from Gmag:')
        vmag = gmag - get_GminusV(teff)
        log.debug(f'  Gmag={gmag:.2f} --> Vmag={vmag:.2f}')
        obs_exp_time = args.get('ExpTime')

        # Force values in to allowed ranges for KPF-etc
        if vmag > 19:
            log.warning(f"Vmag = {vmag:.2f} is outside allowed range for ETC (1-19)")
            vmag = 19
        if vmag < 1:
            log.warning(f"Vmag = {vmag:.2f} is outside allowed range for ETC (1-19)")
            vmag = 1
        if obs_exp_time > 3600:
            log.warning(f"ExpTime = {obs_exp_time:.0f} is outside allowed range for ETC (1-3600)")
            obs_exp_time = 3600
        if obs_exp_time < 1:
            log.warning(f"ExpTime = {obs_exp_time:.0f} is outside allowed range for ETC (1-3600)")
            obs_exp_time = 1
        if teff > 6600:
            log.warning(f"Teff = {teff:.0f} is outside allowed range for ETC (2700-6600)")
            teff = 6600
        if teff < 2700:
            log.warning(f"Teff = {teff:.0f} is outside allowed range for ETC (2700-6600)")
            teff = 2700

        # reference calibration file to scale up/down
        #cal_file = 'KP.20240529.80736.43_L1.fits' # reference etalon L1 file
        data_dir = Path(__file__).parent.parent.parent / 'data'
        cal_file = data_dir / 'KP.20250501.79854.92_L1.fits'

        # Filter wheel populations for both wheels
#         od_arr_scical = [0.1, 0.3, 0.5, 0.8, 1.0, 4.0]
#         od_arr_cal = [0.1, 1.0, 1.3, 2., 3., 4.]
#         od_arr_scical = [0.1, 1.0, 1.3, 2., 3., 4.]
#         od_arr_cal = [0.1, 0.3, 0.5, 0.8, 1.0, 4.0]

        ND1POS = ktl.cache('kpfcal', 'ND1POS')
        ND1POS_allowed_values = list(ND1POS._getEnumerators())
        if 'Unknown' in ND1POS_allowed_values:
            ND1POS_allowed_values.pop(ND1POS_allowed_values.index('Unknown'))
        od_arr_scical = [float(pos[3:]) for pos in ND1POS_allowed_values]

        ND2POS = ktl.cache('kpfcal', 'ND2POS')
        ND2POS_allowed_values = list(ND2POS._getEnumerators())
        if 'Unknown' in ND2POS_allowed_values:
            ND2POS_allowed_values.pop(ND2POS_allowed_values.index('Unknown'))
        od_arr_cal = [float(pos[3:]) for pos in ND2POS_allowed_values]

        od_vals_all, filter_configs_all = all_possible_sums_with_indices_sorted(od_arr_scical, od_arr_cal)

        od, nd_config = get_simulcal_od(vmag, teff, obs_exp_time, cal_file,
                            ref_wave=5500, od_values=od_vals_all,
                            filter_configs=filter_configs_all)

        result = {'CalND1': f'OD {nd_config[0]}',
                  'CalND2': f'OD {nd_config[1]}'}
        log.info(f"Predicted ND1 = {result['CalND1']}")
        log.info(f"Predicted ND2 = {result['CalND2']}")
        tock = datetime.datetime.now()
        elapsed = (tock-tick).total_seconds()
        log.debug(f'ND filter calculation took {elapsed:.1f}s')
        if args.get('set', False):
            SetND.execute(result)
        return result

    @classmethod
    def post_condition(cls, args):
        pass

    @classmethod
    def add_cmdline_args(cls, parser):
        parser.add_argument('Gmag', type=float,
                            help="The gaia G magnitude of the target")
        parser.add_argument('Teff', type=float,
                            help="The effective temperature of the target")
        parser.add_argument('ExpTime', type=float,
                            help="The exposure time on target")
        parser.add_argument("--set", dest="set",
            default=False, action="store_true",
            help="Set these values after calculating?")
        return super().add_cmdline_args(parser)