from itertools import batched
from typing import List
from peewee import JOIN

def query_catalog(sdss_ids: List[int], batch_size: int = 1000):
    """
    Query the SDSS database for photometry and astrometry.
    """

    sdss_ids = sorted(list(set(sdss_ids).difference({0, -1})))

    from almanac.database import is_database_available, catalogdb as db

    for batch in batched(sdss_ids, batch_size):
        q = (
            db
            .SDSS_ID_To_Catalog
            .select(
                db.SDSS_ID_To_Catalog.sdss_id,
                db.SDSS_ID_To_Catalog.catalogid,
                db.SDSS_ID_To_Catalog.version_id,
                db.SDSS_ID_To_Catalog.lead,
                db.SDSS_ID_To_Catalog.allstar_dr17_synspec_rev1,
                db.SDSS_ID_To_Catalog.allwise,
                db.SDSS_ID_To_Catalog.bhm_rm_v0,
                db.SDSS_ID_To_Catalog.bhm_rm_v0_2,
                db.SDSS_ID_To_Catalog.catwise,
                db.SDSS_ID_To_Catalog.catwise2020,
                db.SDSS_ID_To_Catalog.gaia_dr2_source,
                db.SDSS_ID_To_Catalog.gaia_dr3_source,
                db.SDSS_ID_To_Catalog.glimpse,
                db.SDSS_ID_To_Catalog.guvcat,
                db.SDSS_ID_To_Catalog.legacy_survey_dr10,
                db.SDSS_ID_To_Catalog.legacy_survey_dr8,
                db.SDSS_ID_To_Catalog.mangatarget,
                db.SDSS_ID_To_Catalog.marvels_dr11_star,
                db.SDSS_ID_To_Catalog.marvels_dr12_star,
                db.SDSS_ID_To_Catalog.mastar_goodstars,
                db.SDSS_ID_To_Catalog.panstarrs1,
                db.SDSS_ID_To_Catalog.ps1_g18,
                db.SDSS_ID_To_Catalog.sdss_dr13_photoobj,
                db.SDSS_ID_To_Catalog.sdss_dr17_specobj,
                db.SDSS_ID_To_Catalog.skymapper_dr1_1,
                db.SDSS_ID_To_Catalog.skymapper_dr2,
                db.SDSS_ID_To_Catalog.supercosmos,
                db.SDSS_ID_To_Catalog.tic_v8,
                db.SDSS_ID_To_Catalog.twomass_psc,
                db.SDSS_ID_To_Catalog.tycho2,
                db.SDSS_ID_To_Catalog.unwise,
                db.Gaia_DR3.source_id,
                db.Gaia_DR3.solution_id,
                db.Gaia_DR3.designation,
                db.Gaia_DR3.random_index,
                db.Gaia_DR3.ref_epoch,
                db.Gaia_DR3.ra,
                db.Gaia_DR3.ra_error,
                db.Gaia_DR3.dec,
                db.Gaia_DR3.dec_error,
                db.Gaia_DR3.parallax,
                db.Gaia_DR3.parallax_error,
                db.Gaia_DR3.parallax_over_error,
                db.Gaia_DR3.pm,
                db.Gaia_DR3.pmra,
                db.Gaia_DR3.pmra_error,
                db.Gaia_DR3.pmdec,
                db.Gaia_DR3.pmdec_error,
                db.Gaia_DR3.ra_dec_corr,
                db.Gaia_DR3.ra_parallax_corr,
                db.Gaia_DR3.ra_pmra_corr,
                db.Gaia_DR3.ra_pmdec_corr,
                db.Gaia_DR3.dec_parallax_corr,
                db.Gaia_DR3.dec_pmra_corr,
                db.Gaia_DR3.dec_pmdec_corr,
                db.Gaia_DR3.parallax_pmra_corr,
                db.Gaia_DR3.parallax_pmdec_corr,
                db.Gaia_DR3.pmra_pmdec_corr,
                db.Gaia_DR3.astrometric_n_obs_al,
                db.Gaia_DR3.astrometric_n_obs_ac,
                db.Gaia_DR3.astrometric_n_good_obs_al,
                db.Gaia_DR3.astrometric_n_bad_obs_al,
                db.Gaia_DR3.astrometric_gof_al,
                db.Gaia_DR3.astrometric_chi2_al,
                db.Gaia_DR3.astrometric_excess_noise,
                db.Gaia_DR3.astrometric_excess_noise_sig,
                db.Gaia_DR3.astrometric_params_solved,
                db.Gaia_DR3.astrometric_primary_flag,
                db.Gaia_DR3.nu_eff_used_in_astrometry,
                db.Gaia_DR3.pseudocolour,
                db.Gaia_DR3.pseudocolour_error,
                db.Gaia_DR3.ra_pseudocolour_corr,
                db.Gaia_DR3.dec_pseudocolour_corr,
                db.Gaia_DR3.parallax_pseudocolour_corr,
                db.Gaia_DR3.pmra_pseudocolour_corr,
                db.Gaia_DR3.pmdec_pseudocolour_corr,
                db.Gaia_DR3.astrometric_matched_transits,
                db.Gaia_DR3.visibility_periods_used,
                db.Gaia_DR3.astrometric_sigma5d_max,
                db.Gaia_DR3.matched_transits,
                db.Gaia_DR3.new_matched_transits,
                db.Gaia_DR3.matched_transits_removed,
                db.Gaia_DR3.ipd_gof_harmonic_amplitude,
                db.Gaia_DR3.ipd_gof_harmonic_phase,
                db.Gaia_DR3.ipd_frac_multi_peak,
                db.Gaia_DR3.ipd_frac_odd_win,
                db.Gaia_DR3.ruwe,
                db.Gaia_DR3.scan_direction_strength_k1,
                db.Gaia_DR3.scan_direction_strength_k2,
                db.Gaia_DR3.scan_direction_strength_k3,
                db.Gaia_DR3.scan_direction_strength_k4,
                db.Gaia_DR3.scan_direction_mean_k1,
                db.Gaia_DR3.scan_direction_mean_k2,
                db.Gaia_DR3.scan_direction_mean_k3,
                db.Gaia_DR3.scan_direction_mean_k4,
                db.Gaia_DR3.duplicated_source,
                db.Gaia_DR3.phot_g_n_obs,
                db.Gaia_DR3.phot_g_mean_flux,
                db.Gaia_DR3.phot_g_mean_flux_error,
                db.Gaia_DR3.phot_g_mean_flux_over_error,
                db.Gaia_DR3.phot_g_mean_mag,
                db.Gaia_DR3.phot_bp_n_obs,
                db.Gaia_DR3.phot_bp_mean_flux,
                db.Gaia_DR3.phot_bp_mean_flux_error,
                db.Gaia_DR3.phot_bp_mean_flux_over_error,
                db.Gaia_DR3.phot_bp_mean_mag,
                db.Gaia_DR3.phot_rp_n_obs,
                db.Gaia_DR3.phot_rp_mean_flux,
                db.Gaia_DR3.phot_rp_mean_flux_error,
                db.Gaia_DR3.phot_rp_mean_flux_over_error,
                db.Gaia_DR3.phot_rp_mean_mag,
                db.Gaia_DR3.phot_bp_rp_excess_factor,
                db.Gaia_DR3.phot_bp_n_contaminated_transits,
                db.Gaia_DR3.phot_bp_n_blended_transits,
                db.Gaia_DR3.phot_rp_n_contaminated_transits,
                db.Gaia_DR3.phot_rp_n_blended_transits,
                db.Gaia_DR3.phot_proc_mode,
                db.Gaia_DR3.bp_rp,
                db.Gaia_DR3.bp_g,
                db.Gaia_DR3.g_rp,
                db.Gaia_DR3.radial_velocity,
                db.Gaia_DR3.radial_velocity_error,
                db.Gaia_DR3.rv_method_used,
                db.Gaia_DR3.rv_nb_transits,
                db.Gaia_DR3.rv_nb_deblended_transits,
                db.Gaia_DR3.rv_visibility_periods_used,
                db.Gaia_DR3.rv_expected_sig_to_noise,
                db.Gaia_DR3.rv_renormalised_gof,
                db.Gaia_DR3.rv_chisq_pvalue,
                db.Gaia_DR3.rv_time_duration,
                db.Gaia_DR3.rv_amplitude_robust,
                db.Gaia_DR3.rv_template_teff,
                db.Gaia_DR3.rv_template_logg,
                db.Gaia_DR3.rv_template_fe_h,
                db.Gaia_DR3.rv_atm_param_origin,
                db.Gaia_DR3.vbroad,
                db.Gaia_DR3.vbroad_error,
                db.Gaia_DR3.vbroad_nb_transits,
                db.Gaia_DR3.grvs_mag,
                db.Gaia_DR3.grvs_mag_error,
                db.Gaia_DR3.grvs_mag_nb_transits,
                db.Gaia_DR3.rvs_spec_sig_to_noise,
                db.Gaia_DR3.phot_variable_flag,
                db.Gaia_DR3.l,
                db.Gaia_DR3.b,
                db.Gaia_DR3.ecl_lon,
                db.Gaia_DR3.ecl_lat,
                db.Gaia_DR3.in_qso_candidates,
                db.Gaia_DR3.in_galaxy_candidates,
                db.Gaia_DR3.non_single_star,
                db.Gaia_DR3.has_xp_continuous,
                db.Gaia_DR3.has_xp_sampled,
                db.Gaia_DR3.has_rvs,
                db.Gaia_DR3.has_epoch_photometry,
                db.Gaia_DR3.has_epoch_rv,
                db.Gaia_DR3.has_mcmc_gspphot,
                db.Gaia_DR3.has_mcmc_msc,
                db.Gaia_DR3.in_andromeda_survey,
                db.Gaia_DR3.classprob_dsc_combmod_quasar,
                db.Gaia_DR3.classprob_dsc_combmod_galaxy,
                db.Gaia_DR3.classprob_dsc_combmod_star,
                db.Gaia_DR3.teff_gspphot,
                db.Gaia_DR3.teff_gspphot_lower,
                db.Gaia_DR3.teff_gspphot_upper,
                db.Gaia_DR3.logg_gspphot,
                db.Gaia_DR3.logg_gspphot_lower,
                db.Gaia_DR3.logg_gspphot_upper,
                db.Gaia_DR3.mh_gspphot,
                db.Gaia_DR3.mh_gspphot_lower,
                db.Gaia_DR3.mh_gspphot_upper,
                db.Gaia_DR3.distance_gspphot,
                db.Gaia_DR3.distance_gspphot_lower,
                db.Gaia_DR3.distance_gspphot_upper,
                db.Gaia_DR3.azero_gspphot,
                db.Gaia_DR3.azero_gspphot_lower,
                db.Gaia_DR3.azero_gspphot_upper,
                db.Gaia_DR3.ag_gspphot,
                db.Gaia_DR3.ag_gspphot_lower,
                db.Gaia_DR3.ag_gspphot_upper,
                db.Gaia_DR3.ebpminrp_gspphot,
                db.Gaia_DR3.ebpminrp_gspphot_lower,
                db.Gaia_DR3.ebpminrp_gspphot_upper,
                db.Gaia_DR3.libname_gspphot,
                db.TwoMassPSC.designation,
                db.TwoMassPSC.j_m,
                db.TwoMassPSC.j_cmsig,
                db.TwoMassPSC.j_msigcom,
                db.TwoMassPSC.j_snr,
                db.TwoMassPSC.h_m,
                db.TwoMassPSC.h_cmsig,
                db.TwoMassPSC.h_msigcom,
                db.TwoMassPSC.h_snr,
                db.TwoMassPSC.k_m,
                db.TwoMassPSC.k_cmsig,
                db.TwoMassPSC.k_msigcom,
                db.TwoMassPSC.k_snr,
                db.TwoMassPSC.ph_qual,
                db.TwoMassPSC.rd_flg,
                db.TwoMassPSC.bl_flg,
                db.TwoMassPSC.cc_flg,
                db.TwoMassPSC.ndet,
                db.TwoMassPSC.prox,
                db.TwoMassPSC.pxpa,
                db.TwoMassPSC.pxcntr,
                db.TwoMassPSC.gal_contam,
                db.TwoMassPSC.mp_flg,
                db.TwoMassPSC.scan,
                db.TwoMassPSC.glon,
                db.TwoMassPSC.glat,
                db.TwoMassPSC.x_scan,
                db.TwoMassPSC.jdate,
                db.TwoMassPSC.j_psfchi,
                db.TwoMassPSC.h_psfchi,
                db.TwoMassPSC.k_psfchi,
                db.TwoMassPSC.j_m_stdap,
                db.TwoMassPSC.j_msig_stdap,
                db.TwoMassPSC.h_m_stdap,
                db.TwoMassPSC.h_msig_stdap,
                db.TwoMassPSC.k_m_stdap,
                db.TwoMassPSC.k_msig_stdap,
                db.TwoMassPSC.dist_edge_ns,
                db.TwoMassPSC.dist_edge_ew,
                db.TwoMassPSC.dist_edge_flg,
                db.TwoMassPSC.dup_src,
                db.TwoMassPSC.use_src,
                db.TwoMassPSC.a,
                db.TwoMassPSC.dist_opt,
                db.TwoMassPSC.phi_opt,
                db.TwoMassPSC.b_m_opt,
                db.TwoMassPSC.vr_m_opt,
                db.TwoMassPSC.nopt_mchs,
                db.TwoMassPSC.ext_key,
                db.TwoMassPSC.scan_key,
                db.TwoMassPSC.coadd_key,
                db.TwoMassPSC.coadd
            )
            .distinct(db.SDSS_ID_To_Catalog.sdss_id)
            .join(db.Gaia_DR3, join_type=JOIN.LEFT_OUTER)
            .switch(db.SDSS_ID_To_Catalog)
            .join(db.TwoMassPSC, join_type=JOIN.LEFT_OUTER)
            .where(db.SDSS_ID_To_Catalog.sdss_id.in_(tuple(batch)))
            .dicts()
        )
        yield from q.iterator()
