#!/usr/bin/env python
"""
send in DES meds files for this
"""

import fitcosmos
import argparse
import meds
import yaml
import fitsio
import numpy as np
import esutil as eu

parser=argparse.ArgumentParser()
parser.add_argument('--conf',required=True)
parser.add_argument('--plot',required=True)
parser.add_argument('--output',required=True)
parser.add_argument('meds')

parser.add_argument('--extra-psf-fwhm',type=float,default=0.0)

parser.add_argument('--plot-only',action='store_true')

FWHM_FAC = 2*np.sqrt(2*np.log(2))

def main():
    args=parser.parse_args()
    assert 'meds' not in args.output

    with open(args.conf) as fobj:
        conf = yaml.load(fobj)
        fof_conf = conf['fofs']

    print('loading:',args.meds)
    #assert 'des' in args.meds.lower(),'send only DES meds for this task'

    m = meds.MEDS(args.meds)

    if args.plot_only:
        fofs=fitsio.read(args.output, ext='fofs')
    else:
        # should be in arcsec
        radcol = fof_conf['radius_column']
        assert 'arcsec' in radcol,'expected units of arcsec'
        rad = m[radcol]

        # should be in arcsec
        psf_sigma = args.extra_psf_fwhm/FWHM_FAC

        rad = np.sqrt(rad**2 + psf_sigma**2)

        m._cat[radcol] = rad

        print('getting fof groups')
        nbr_data, fofs = fitcosmos.fofs.get_fofs(
            m,
            fof_conf,
        )
        print('writing:',args.output)
        with fitsio.FITS(args.output,'rw',clobber=True) as fits:
            fits.write(fofs,extname='fofs')
            fits.write(nbr_data,extname='nbrs')

    fitcosmos.fofs.plot_fofs(
        m,
        fofs,
        plotfile=args.plot,
        width=2000,
        fof_type='filled circle',
        fof_size=0.2,
    )

if __name__=='__main__':
    main()
 
