"""Extract the columns needed from the GALAH DR4 / DR3 allstar FITS files into compact npz caches."""
import numpy as np
from astropy.io import fits
from common import RAW, CACHE, NCAP


def native(a):
    a = np.array(a)
    return a.astype(a.dtype.newbyteorder('=')) if a.dtype.byteorder == '>' else a

def prep_dr4():
    h = fits.open(RAW + 'galah_dr4_allstar_240705.fits', memmap=True)
    d = h[1].data
    out = {}
    for c in ['sobject_id', 'gaiadr3_source_id', 'ra', 'dec', 'flag_sp', 'flag_red', 'snr_px_ccd1',
              'snr_px_ccd2', 'snr_px_ccd3', 'snr_px_ccd4', 'chi2_sp', 'teff', 'logg', 'fe_h',
              'e_fe_h', 'vmic', 'vsini', 'rv_comp_1', 'rv_comp_2', 'rv_gaia_dr3', 'mass', 'age',
              'r_med', 'phot_g_mean_mag', 'bp_rp', 'ruwe', 'parallax', 'ebv', 'ks_m', 'W2mag',
              'ew_h_alpha', 'c_fe', 'flag_c_fe', 'mg_fe', 'ca_fe', 'sc_fe', 'ti_fe', 'ni_fe',
              'flag_sc_fe', 'flag_ca_fe', 'cr_fe', 'flag_cr_fe']:
        out[c] = native(d[c])
    out['tmass_id'] = np.array(d['tmass_id']).astype('U17')
    out['survey_name'] = np.array(d['survey_name']).astype('U20')
    for e in NCAP:
        out[e + '_fe'] = np.array(d[e + '_fe'], dtype='f4')
        out['e_' + e + '_fe'] = np.array(d['e_' + e + '_fe'], dtype='f4')
        out['flag_' + e + '_fe'] = np.array(d['flag_' + e + '_fe'], dtype='i2')
    np.savez(CACHE + 'dr4.npz', **out)
    print('DR4 rows', len(d))

def prep_dr3():
    h = fits.open(RAW + 'GALAH_DR3_main_allstar_v2.fits', memmap=True)
    d = h[1].data
    out = {}
    for c in ['sobject_id', 'dr3_source_id', 'flag_sp', 'teff', 'logg', 'fe_h', 'flag_fe_h',
              'vbroad', 'chi2_sp', 'snr_c1_iraf', 'snr_c2_iraf', 'snr_c3_iraf', 'ra_dr2',
              'dec_dr2', 'flag_repeat']:
        out[c] = native(d[c])
    out['star_id'] = np.array(d['star_id']).astype('U17')
    for e in NCAP:
        E = e.capitalize()
        out[e + '_fe'] = np.array(d[E + '_fe'], dtype='f4')
        out['e_' + e + '_fe'] = np.array(d['e_' + E + '_fe'], dtype='f4')
        out['flag_' + e + '_fe'] = np.array(d['flag_' + E + '_fe'], dtype='i2')
        out['nr_' + e + '_fe'] = np.array(d['nr_' + E + '_fe'], dtype='f4')
    np.savez(CACHE + 'dr3.npz', **out)
    print('DR3 rows', len(d))

def main():
    prep_dr4()
    prep_dr3()


if __name__ == '__main__':
    main()