#!/usr/bin/env python

import astropy.io.fits as fits
import numpy as np
import os
import matplotlib.pyplot as plt
import argparse

def _fname(expid, night, 
           basedir='/n/home/datasystems/users/ameisner/reduced/focus',
           ccds=False):

           fname = basedir + '/' + night + '/' + str(expid).zfill(8) + '/gfa-' + str(expid).zfill(8) + '_psfs.fits'    

           if ccds:
               fname = fname.replace('_psfs.fits', '_ccds.fits')

           return fname

def _actual_expid_list(expids, night, basedir='/n/home/datasystems/users/ameisner/reduced/focus'):

    keep = []
    for i, expid in enumerate(expids):
        fname = _fname(expid, night, basedir=basedir, ccds=True)
        if not os.path.exists(fname):
            continue
        tab = fits.getdata(fname)
        # try to handle case where observer accidentally lists the 'setup focus scan'
        # 1 second exposure as the start of the focus scan
        if (i == 0) & (tab[0]['EXPTIME'] < 1.1):
            print('SKIPPING DUMMY SETUP EXPOSURE')
            continue

        program = tab[0]['PROGRAM'].strip()

        if program != 'focus scan':
            break

        keep.append(expid)

    return keep

def focus_plots(night, expids,
               basedir='/n/home/datasystems/users/ameisner/reduced/focus',
               outdir='/n/home/desiobserver/focus_scan_pngs', no_popups=False):

    expids = _actual_expid_list(expids, night, basedir=basedir)

    if len(expids) == 0:
        print('NO FOCUS SCAN EXPOSURES TO ANALYZE ??')
        assert(False)

    plt.figure(1, figsize=(12.0*(len(expids)/7.0), 9))
    extnames = ['GUIDE0', 'GUIDE2', 'GUIDE3', 'GUIDE5', 'GUIDE7', 'GUIDE8']

    focus_z = []
    fwhm_pix = []

    # PSF stamps plot
    plt.subplots_adjust(hspace=0.01, wspace=0.01)
    for i, expid in enumerate(expids):
        fname = _fname(expid, night, basedir=basedir)
        print(fname)
        fname_ccds = _fname(expid, night, basedir=basedir, ccds=True)
        if not os.path.exists(fname):
            continue
        ccds = fits.getdata(fname_ccds)

        if np.sum(np.isfinite(ccds['PSF_FWHM_PIX'])) != 0:
            fwhm_pix.append(np.median(ccds['PSF_FWHM_PIX'][np.isfinite(ccds['PSF_FWHM_PIX'])]))
            focus_z.append(float(ccds[0]['FOCUS'].split(',')[2]))
            
        hdul = fits.open(fname)
        extnames_present = [hdu.header['EXTNAME'] for hdu in hdul]
        for j, extname in enumerate(extnames):
            if extname not in extnames_present:
                continue
            print(i, j)
            plt.subplot(6, len(expids), len(expids)*j + i +  1)
            plt.xticks([])
            plt.yticks([])
            im = fits.getdata(fname, extname=extname)
            plt.imshow(im, interpolation='nearest', origin='lower', cmap='gray_r', vmin=0.01)
            plt.text(5, 44, str(expid) + '; ' + extname, color='r', fontsize=9)
            plt.text(10, 3.5, 'z = ' + str(int(float(ccds[0]['FOCUS'].split(',')[2]))), color='r')
            
            if np.isfinite(ccds[j]['XCENTROID_PSF']) and np.isfinite(ccds[j]['YCENTROID_PSF']):
                plt.scatter([ccds[j]['XCENTROID_PSF']], [ccds[j]['YCENTROID_PSF']], marker='.', c='r')

    
    expid_min = int(np.min(expids))

    print(focus_z)
    print(fwhm_pix)

    plt.savefig(os.path.join(outdir, 'stamps_focus_scan-' + str(expid_min).zfill(8)+'.png'), bbox_inches='tight')
    #plt.cla()

    plt.figure(200)
    
    asec_per_pix = 0.205

    focus_z = np.array(focus_z)
    fwhm_asec = np.array(fwhm_pix)*asec_per_pix
    plt.scatter(focus_z, fwhm_asec)
    plt.xlabel('focus z (micron)')
    plt.ylabel('FWHM (asec)')

    coeff = np.polyfit(focus_z, fwhm_asec, 2)

    xsamp = np.arange(np.min(focus_z), np.max(focus_z))
    ysamp = coeff[0]*(np.power(xsamp, 2)) + coeff[1]*xsamp + coeff[2]

    plt.title('focus scan starting with EXPID = ' + str(expid_min))

    
    plt.plot(xsamp, ysamp)

    zmin = -coeff[1]/(2*coeff[0])

    min_fwhm_fit_asec = coeff[0]*(zmin**2) + coeff[1]*zmin + coeff[2]
    
    yrange = [np.min(fwhm_asec), np.max(fwhm_asec)]
    plt.text(focus_z[2], yrange[0] + 0.8*(yrange[1]-yrange[0]), 'best FWHM (meas) : ' + '{:.2f}'.format(np.min(fwhm_asec)))
    plt.text(focus_z[2], yrange[0] + 0.7*(yrange[1]-yrange[0]), 'best FWHM (fit) : ' + '{:.2f}'.format(min_fwhm_fit_asec))
    plt.text(focus_z[2], yrange[0] + 0.9*(yrange[1]-yrange[0]), 'best focus : ' + str(int(np.round(zmin))))
    
    plt.savefig(os.path.join(outdir, 'fit_focus_scan-' + str(expid_min).zfill(8) + '.png'), bbox_inches='tight')
    if not no_popups:
        plt.show()
    
def _test():
    night = '20200131'
    expids = 45446 + np.arange(7)

    focus_plots(night, expids, basedir='/project/projectdirs/desi/users/ameisner/GFA/run/psf_flux_weighted_centroid', outdir='.')

def _test_missing_cam():
    night = '20200131'
    expids = 45485 + np.arange(7)

    focus_plots(night, expids, basedir='/project/projectdirs/desi/users/ameisner/GFA/run/psf_flux_weighted_centroid')

if __name__ == "__main__":
    descr = 'GFA focus sequence plots/analysis'
    parser = argparse.ArgumentParser(description=descr)

    parser.add_argument('first_expid', type=int, nargs=1)

    parser.add_argument('night', type=str, nargs=1)
    
    parser.add_argument('--basedir', default='/n/home/datasystems/users/ameisner/reduced/focus',
                        type=str, help='base directory for GFA reductions')

    parser.add_argument('--outdir', default='/n/home/desiobserver/focus_scan_pngs', 
                        type=str, help='output directory for plot PNGs')

    parser.add_argument('--no_popups', default=False, action='store_true',
                        help='write PNGs without popping up plot windows')

    args = parser.parse_args()

    expids = args.first_expid + np.arange(16, dtype=int)

    print(expids)
    print(args.night[0])
    print(args.basedir)
    
    outdir = args.outdir if os.path.exists(args.outdir) else '.'
    focus_plots(args.night[0], expids, basedir=args.basedir, outdir=outdir, no_popups=args.no_popups)

[((1758, 1803), 'matplotlib.pyplot.subplots_adjust', 'plt.subplots_adjust', ([], {'hspace': '(0.01)', 'wspace': '(0.01)'}), '(hspace=0.01, wspace=0.01)\n', (1777, 1803), True, 'import matplotlib.pyplot as plt\n'), ((3439, 3454), 'matplotlib.pyplot.figure', 'plt.figure', (['(200)'], {}), '(200)\n', (3449, 3454), True, 'import matplotlib.pyplot as plt\n'), ((3500, 3517), 'numpy.array', 'np.array', (['focus_z'], {}), '(focus_z)\n', (3508, 3517), True, 'import numpy as np\n'), ((3570, 3601), 'matplotlib.pyplot.scatter', 'plt.scatter', (['focus_z', 'fwhm_asec'], {}), '(focus_z, fwhm_asec)\n', (3581, 3601), True, 'import matplotlib.pyplot as plt\n'), ((3606, 3636), 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""focus z (micron)"""'], {}), "('focus z (micron)')\n", (3616, 3636), True, 'import matplotlib.pyplot as plt\n'), ((3641, 3666), 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""FWHM (asec)"""'], {}), "('FWHM (asec)')\n", (3651, 3666), True, 'import matplotlib.pyplot as plt\n'), ((3680, 3713), 'numpy.polyfit', 'np.polyfit', (['focus_z', 'fwhm_asec', '(2)'], {}), '(focus_z, fwhm_asec, 2)\n', (3690, 3713), True, 'import numpy as np\n'), ((3920, 3942), 'matplotlib.pyplot.plot', 'plt.plot', (['xsamp', 'ysamp'], {}), '(xsamp, ysamp)\n', (3928, 3942), True, 'import matplotlib.pyplot as plt\n'), ((5117, 5159), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': 'descr'}), '(description=descr)\n', (5140, 5159), False, 'import argparse\n'), ((777, 796), 'astropy.io.fits.getdata', 'fits.getdata', (['fname'], {}), '(fname)\n', (789, 796), True, 'import astropy.io.fits as fits\n'), ((2062, 2086), 'astropy.io.fits.getdata', 'fits.getdata', (['fname_ccds'], {}), '(fname_ccds)\n', (2074, 2086), True, 'import astropy.io.fits as fits\n'), ((2337, 2353), 'astropy.io.fits.open', 'fits.open', (['fname'], {}), '(fname)\n', (2346, 2353), True, 'import astropy.io.fits as fits\n'), ((3248, 3262), 'numpy.min', 'np.min', (['expids'], {}), '(expids)\n', (3254, 3262), True, 'import numpy as np\n'), ((3534, 3552), 'numpy.array', 'np.array', (['fwhm_pix'], {}), '(fwhm_pix)\n', (3542, 3552), True, 'import numpy as np\n'), ((3737, 3752), 'numpy.min', 'np.min', (['focus_z'], {}), '(focus_z)\n', (3743, 3752), True, 'import numpy as np\n'), ((3754, 3769), 'numpy.max', 'np.max', (['focus_z'], {}), '(focus_z)\n', (3760, 3769), True, 'import numpy as np\n'), ((4068, 4085), 'numpy.min', 'np.min', (['fwhm_asec'], {}), '(fwhm_asec)\n', (4074, 4085), True, 'import numpy as np\n'), ((4087, 4104), 'numpy.max', 'np.max', (['fwhm_asec'], {}), '(fwhm_asec)\n', (4093, 4104), True, 'import numpy as np\n'), ((4609, 4619), 'matplotlib.pyplot.show', 'plt.show', ([], {}), '()\n', (4617, 4619), True, 'import matplotlib.pyplot as plt\n'), ((4682, 4694), 'numpy.arange', 'np.arange', (['(7)'], {}), '(7)\n', (4691, 4694), True, 'import numpy as np\n'), ((4896, 4908), 'numpy.arange', 'np.arange', (['(7)'], {}), '(7)\n', (4905, 4908), True, 'import numpy as np\n'), ((5825, 5849), 'numpy.arange', 'np.arange', (['(16)'], {'dtype': 'int'}), '(16, dtype=int)\n', (5834, 5849), True, 'import numpy as np\n'), ((5951, 5978), 'os.path.exists', 'os.path.exists', (['args.outdir'], {}), '(args.outdir)\n', (5965, 5978), False, 'import os\n'), ((719, 740), 'os.path.exists', 'os.path.exists', (['fname'], {}), '(fname)\n', (733, 740), False, 'import os\n'), ((2003, 2024), 'os.path.exists', 'os.path.exists', (['fname'], {}), '(fname)\n', (2017, 2024), False, 'import os\n'), ((2641, 2655), 'matplotlib.pyplot.xticks', 'plt.xticks', (['[]'], {}), '([])\n', (2651, 2655), True, 'import matplotlib.pyplot as plt\n'), ((2668, 2682), 'matplotlib.pyplot.yticks', 'plt.yticks', (['[]'], {}), '([])\n', (2678, 2682), True, 'import matplotlib.pyplot as plt\n'), ((2700, 2736), 'astropy.io.fits.getdata', 'fits.getdata', (['fname'], {'extname': 'extname'}), '(fname, extname=extname)\n', (2712, 2736), True, 'import astropy.io.fits as fits\n'), ((2749, 2835), 'matplotlib.pyplot.imshow', 'plt.imshow', (['im'], {'interpolation': '"""nearest"""', 'origin': '"""lower"""', 'cmap': '"""gray_r"""', 'vmin': '(0.01)'}), "(im, interpolation='nearest', origin='lower', cmap='gray_r', vmin\n    =0.01)\n", (2759, 2835), True, 'import matplotlib.pyplot as plt\n'), ((2106, 2139), 'numpy.isfinite', 'np.isfinite', (["ccds['PSF_FWHM_PIX']"], {}), "(ccds['PSF_FWHM_PIX'])\n", (2117, 2139), True, 'import numpy as np\n'), ((3038, 3075), 'numpy.isfinite', 'np.isfinite', (["ccds[j]['XCENTROID_PSF']"], {}), "(ccds[j]['XCENTROID_PSF'])\n", (3049, 3075), True, 'import numpy as np\n'), ((3080, 3117), 'numpy.isfinite', 'np.isfinite', (["ccds[j]['YCENTROID_PSF']"], {}), "(ccds[j]['YCENTROID_PSF'])\n", (3091, 3117), True, 'import numpy as np\n'), ((3135, 3226), 'matplotlib.pyplot.scatter', 'plt.scatter', (["[ccds[j]['XCENTROID_PSF']]", "[ccds[j]['YCENTROID_PSF']]"], {'marker': '"""."""', 'c': '"""r"""'}), "([ccds[j]['XCENTROID_PSF']], [ccds[j]['YCENTROID_PSF']], marker=\n    '.', c='r')\n", (3146, 3226), True, 'import matplotlib.pyplot as plt\n'), ((3793, 3811), 'numpy.power', 'np.power', (['xsamp', '(2)'], {}), '(xsamp, 2)\n', (3801, 3811), True, 'import numpy as np\n'), ((4210, 4227), 'numpy.min', 'np.min', (['fwhm_asec'], {}), '(fwhm_asec)\n', (4216, 4227), True, 'import numpy as np\n'), ((4443, 4457), 'numpy.round', 'np.round', (['zmin'], {}), '(zmin)\n', (4451, 4457), True, 'import numpy as np\n'), ((2206, 2239), 'numpy.isfinite', 'np.isfinite', (["ccds['PSF_FWHM_PIX']"], {}), "(ccds['PSF_FWHM_PIX'])\n", (2217, 2239), True, 'import numpy as np\n')]