import astropy.io.fits as fits
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import obj_data as od
import saphires as saph
from astropy.time import Time
from astropy.visualization import ZScaleInterval, SqrtStretch, ImageNormalize
from matplotlib.backends.backend_pdf import PdfPages

ra = od.ra
dec = od.dec
pmra = od.pmra
pmdec = od.pmdec
plx  = od.plx
epoch = od.epoch
matplotlib.rcParams.update({'font.size': 12})

def write_fits(fn, data, im_headers, wcs_header):
    '''
    Writes a new fits file including the image data and
    and updated header for the new image

        Parameters
        ----------
        fn: string
            The desired file name of the new fits file
        data: array-like
            Contains all the image data

        Returns
        -------
        avg_airmass: float
            the amount of atmosphere obscuring the target, found in image header. Here
            the airmass for all images is averaged
        bjd: float
            Barycentric Julian Date, found in the image header
        header: Header
    '''
    for keys in wcs_header:
        if keys not in ['HISTORY', 'COMMENT']:
            im_headers[0][keys] = wcs_header[keys]
    airmass = []
    for i in im_headers:
        airmass.append(i['AIRMASS'])
    avg_airmass = np.mean(airmass)
    im_headers[0]['AIRMASS'] = avg_airmass
    jd_middle = np.zeros(len(im_headers))
    for i in range(len(im_headers)):
        jd_middle[i] = Time(im_headers[i]['DATE-OBS'], format='isot').jd
        exptime = im_headers[i]['EXPTIME']
        jd_middle[i] = jd_middle[i] + (exptime/2.0)/3600.0/24.0
    isot_date_obs = Time(np.mean(jd_middle), format='jd').isot
    tele = im_headers[0]['SITEID']
    brv,bjd,bvcorr = saph.utils.brvc(isot_date_obs,0.0,tele,ra=ra,dec=dec,epoch=epoch, pmra=pmra, pmdec=pmdec, px=plx)
    im_headers[0]['BJD'] = bjd[0]
    header = im_headers[0]
    hdu_p = fits.PrimaryHDU(data=data, header=header)
    hdu = fits.HDUList([hdu_p])
    hdu.writeto(fn)

    return avg_airmass, bjd, header

def write_pdf(name, images, model=None, final_stars=None, residual_stars=None, fluxes=None, plot_res=None):
    pp = PdfPages(name)
    for i in range(len(images)):
        fig, ax = plt.subplots(1, figsize=(10, 10))
        norm = ImageNormalize(images[i], interval=ZScaleInterval(), stretch=SqrtStretch())
        im = ax.imshow(images[i], norm=norm)
        plt.colorbar(im)
        plt.tight_layout()
        pp.savefig()
        plt.close()
    if model is not None:
        fig, ax = plt.subplots(1, figsize=(10, 10))
        psf = ax.imshow(model)
        plt.colorbar(psf)
        ax.set_title('PSF Model')
        plt.tight_layout()
        pp.savefig()
        plt.close()
    if final_stars is not None:
        if plot_res == 'y':
            nrows = len(final_stars)
            ncols = 2
            fig, ax = plt.subplots(nrows=nrows, ncols=ncols, figsize=(10, 800), squeeze=True)
            ax = ax.ravel()
            index = 0
            for i in range(0, nrows*ncols, 2):
                norm = simple_norm(final_stars[index],'log')
                norm2 = simple_norm(residual_stars[index], 'linear')
                im = ax[i].imshow(final_stars[index], norm=norm, origin='lower', cmap='viridis', interpolation='none')
                fig.colorbar(im, ax = ax[i])
                ax[i].set_title(np.str(fluxes[index]))
                im_r = ax[i+1].imshow(residual_stars[index], norm=norm2, origin='lower', cmap='viridis', interpolation='none')
                fig.colorbar(im_r, ax = ax[i+1])
                index = index + 1
            plt.tight_layout()
            pp.savefig()
            plt.close()
    pp.close()

def write_csv(name, im_name, bjd, filt, airmass, results, sky):
    f = open(name, 'w')
    f.write('NAME, ID, BJD, FLUX, FLUX ERROR, MAG, MAG ERROR, FILTER, X POSITION, Y POSITION, AIRMASS, RA, DEC\n')
    for i in range(sky.size):
        if results['flux_fit'][i] > 0:
            star_id = results['id'][i]
            flux = results['flux_fit'][i]
            fluxerr = results['flux_unc'][i]
            mag = -2.5*np.log10(flux)
            magerr = (1.08574*fluxerr)/(flux)
            x_pos = results['x_fit'][i]
            y_pos = results['y_fit'][i]
            ra = sky[i].ra.degree
            dec = sky[i].dec.degree
            f.write(im_name+','+np.str(i)+','+np.str(bjd)+','+np.str(flux)+','+np.str(fluxerr)+','+np.str(mag)+','+np.str(magerr)
                    +','+filt+','+np.str(x_pos)+','+np.str(y_pos)+','+str(airmass)+','+np.str(ra)+','+np.str(dec)+'\n')
    f.close()

def write_txt(name, sources, stars_tbl, fwhm, results=None, t0=None,t1=None,t2=None,t3=None,t4=None,t5=None):
    '''
    Short text file with diagnostic info about each image set, specifically
    for a successful run of the image set

        Parameters
        ----------
        name: string
            name of the saved file
        sources: Table
            tabulated info about all the stars found on the image
        stars_tbl: Table
            tabulated info about all the stars used to form a psf
        results: Table
            tabulated info about all the stars found with the photometry routine

    '''
    f = open(name, 'w')
    f.write('Number of stars in sources: '+np.str(len(sources))+'\nNumber of stars in stars_tbl: '+np.str(len(stars_tbl))
            +'\nNumbers of stars in results: '+np.str(len(results))+'\nMin, Max, Median peaks in sources: '
            +np.str(np.min(sources['peak']))+', '+np.str(np.max(sources['peak']))+', '+np.str(np.median(sources['peak']))
            +'\nMin, Max, Median fluxes in results: '+np.str(np.min(results['flux_fit']))+', '+np.str(np.max(results['flux_fit']))+', '
            +np.str(np.median(results['flux_fit']))+'\nFWHM: '+np.str(fwhm)+'\n')
    if t5:
        t_1 = t1-t0
        t_2 = t2-t1
        t_3 = t3-t2
        t_4 = t4-t3
        t_5 = t5-t4
        t_f = t5-t0
        f.write('Time to combine images: '+np.str(t_1)+'\nTime to find stars: '+np.str(t_2)+'\nTime to build psf: '
                +np.str(t_3)+'\nTime to run photometry: '+np.str(t_4)+'\nTime to get wcs: '+np.str(t_5)+'\nTotal time: '
                +np.str(t_f)+'\n')
    f.close()

[((17, 0, 17, 45), 'matplotlib.rcParams.update', 'matplotlib.rcParams.update', ({(17, 27, 17, 44): "{'font.size': 12}"}, {}), "({'font.size': 12})", False, 'import matplotlib\n'), ((46, 18, 46, 34), 'numpy.mean', 'np.mean', ({(46, 26, 46, 33): 'airmass'}, {}), '(airmass)', True, 'import numpy as np\n'), ((55, 21, 55, 118), 'saphires.utils.brvc', 'saph.utils.brvc', (), '', True, 'import saphires as saph\n'), ((58, 12, 58, 53), 'astropy.io.fits.PrimaryHDU', 'fits.PrimaryHDU', (), '', True, 'import astropy.io.fits as fits\n'), ((59, 10, 59, 31), 'astropy.io.fits.HDUList', 'fits.HDUList', ({(59, 23, 59, 30): '[hdu_p]'}, {}), '([hdu_p])', True, 'import astropy.io.fits as fits\n'), ((65, 9, 65, 23), 'matplotlib.backends.backend_pdf.PdfPages', 'PdfPages', ({(65, 18, 65, 22): 'name'}, {}), '(name)', False, 'from matplotlib.backends.backend_pdf import PdfPages\n'), ((67, 18, 67, 51), 'matplotlib.pyplot.subplots', 'plt.subplots', (), '', True, 'import matplotlib.pyplot as plt\n'), ((70, 8, 70, 24), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ({(70, 21, 70, 23): 'im'}, {}), '(im)', True, 'import matplotlib.pyplot as plt\n'), ((71, 8, 71, 26), 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((73, 8, 73, 19), 'matplotlib.pyplot.close', 'plt.close', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((75, 18, 75, 51), 'matplotlib.pyplot.subplots', 'plt.subplots', (), '', True, 'import matplotlib.pyplot as plt\n'), ((77, 8, 77, 25), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ({(77, 21, 77, 24): 'psf'}, {}), '(psf)', True, 'import matplotlib.pyplot as plt\n'), ((79, 8, 79, 26), 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((81, 8, 81, 19), 'matplotlib.pyplot.close', 'plt.close', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((50, 23, 50, 69), 'astropy.time.Time', 'Time', (), '', False, 'from astropy.time import Time\n'), ((53, 25, 53, 43), 'numpy.mean', 'np.mean', ({(53, 33, 53, 42): 'jd_middle'}, {}), '(jd_middle)', True, 'import numpy as np\n'), ((86, 22, 86, 93), 'matplotlib.pyplot.subplots', 'plt.subplots', (), '', True, 'import matplotlib.pyplot as plt\n'), ((98, 12, 98, 30), 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((100, 12, 100, 23), 'matplotlib.pyplot.close', 'plt.close', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((68, 50, 68, 66), 'astropy.visualization.ZScaleInterval', 'ZScaleInterval', ({}, {}), '()', False, 'from astropy.visualization import ZScaleInterval, SqrtStretch, ImageNormalize\n'), ((68, 76, 68, 89), 'astropy.visualization.SqrtStretch', 'SqrtStretch', ({}, {}), '()', False, 'from astropy.visualization import ZScaleInterval, SqrtStretch, ImageNormalize\n'), ((111, 23, 111, 37), 'numpy.log10', 'np.log10', ({(111, 32, 111, 36): 'flux'}, {}), '(flux)', True, 'import numpy as np\n'), ((143, 63, 143, 75), 'numpy.str', 'np.str', ({(143, 70, 143, 74): 'fwhm'}, {}), '(fwhm)', True, 'import numpy as np\n'), ((94, 32, 94, 53), 'numpy.str', 'np.str', ({(94, 39, 94, 52): 'fluxes[index]'}, {}), '(fluxes[index])', True, 'import numpy as np\n'), ((153, 17, 153, 28), 'numpy.str', 'np.str', ({(153, 24, 153, 27): 't_f'}, {}), '(t_f)', True, 'import numpy as np\n'), ((118, 102, 118, 113), 'numpy.str', 'np.str', ({(118, 109, 118, 112): 'dec'}, {}), '(dec)', True, 'import numpy as np\n'), ((143, 20, 143, 50), 'numpy.median', 'np.median', ({(143, 30, 143, 49): "results['flux_fit']"}, {}), "(results['flux_fit'])", True, 'import numpy as np\n'), ((152, 92, 152, 103), 'numpy.str', 'np.str', ({(152, 99, 152, 102): 't_5'}, {}), '(t_5)', True, 'import numpy as np\n'), ((118, 87, 118, 97), 'numpy.str', 'np.str', ({(118, 94, 118, 96): 'ra'}, {}), '(ra)', True, 'import numpy as np\n'), ((142, 102, 142, 129), 'numpy.max', 'np.max', ({(142, 109, 142, 128): "results['flux_fit']"}, {}), "(results['flux_fit'])", True, 'import numpy as np\n'), ((152, 58, 152, 69), 'numpy.str', 'np.str', ({(152, 65, 152, 68): 't_4'}, {}), '(t_4)', True, 'import numpy as np\n'), ((142, 61, 142, 88), 'numpy.min', 'np.min', ({(142, 68, 142, 87): "results['flux_fit']"}, {}), "(results['flux_fit'])", True, 'import numpy as np\n'), ((152, 17, 152, 28), 'numpy.str', 'np.str', ({(152, 24, 152, 27): 't_3'}, {}), '(t_3)', True, 'import numpy as np\n'), ((118, 52, 118, 65), 'numpy.str', 'np.str', ({(118, 59, 118, 64): 'y_pos'}, {}), '(y_pos)', True, 'import numpy as np\n'), ((141, 94, 141, 120), 'numpy.median', 'np.median', ({(141, 104, 141, 119): "sources['peak']"}, {}), "(sources['peak'])", True, 'import numpy as np\n'), ((151, 80, 151, 91), 'numpy.str', 'np.str', ({(151, 87, 151, 90): 't_2'}, {}), '(t_2)', True, 'import numpy as np\n'), ((118, 34, 118, 47), 'numpy.str', 'np.str', ({(118, 41, 118, 46): 'x_pos'}, {}), '(x_pos)', True, 'import numpy as np\n'), ((141, 57, 141, 80), 'numpy.max', 'np.max', ({(141, 64, 141, 79): "sources['peak']"}, {}), "(sources['peak'])", True, 'import numpy as np\n'), ((151, 43, 151, 54), 'numpy.str', 'np.str', ({(151, 50, 151, 53): 't_1'}, {}), '(t_1)', True, 'import numpy as np\n'), ((141, 20, 141, 43), 'numpy.min', 'np.min', ({(141, 27, 141, 42): "sources['peak']"}, {}), "(sources['peak'])", True, 'import numpy as np\n'), ((117, 115, 117, 129), 'numpy.str', 'np.str', ({(117, 122, 117, 128): 'magerr'}, {}), '(magerr)', True, 'import numpy as np\n'), ((117, 99, 117, 110), 'numpy.str', 'np.str', ({(117, 106, 117, 109): 'mag'}, {}), '(mag)', True, 'import numpy as np\n'), ((117, 79, 117, 94), 'numpy.str', 'np.str', ({(117, 86, 117, 93): 'fluxerr'}, {}), '(fluxerr)', True, 'import numpy as np\n'), ((117, 62, 117, 74), 'numpy.str', 'np.str', ({(117, 69, 117, 73): 'flux'}, {}), '(flux)', True, 'import numpy as np\n'), ((117, 46, 117, 57), 'numpy.str', 'np.str', ({(117, 53, 117, 56): 'bjd'}, {}), '(bjd)', True, 'import numpy as np\n'), ((117, 32, 117, 41), 'numpy.str', 'np.str', ({(117, 39, 117, 40): 'i'}, {}), '(i)', True, 'import numpy as np\n')]