#!/usr/bin/env python
"""
Create stacked image from a subsample of a source list.
"""

import os
import sys
import cPickle as pickle

import Ska.Table
import Ska.Numpy
from os.path import join
import pyfits
import numpy as np
import yaml

def get_options():
    from optparse import OptionParser
    parser = OptionParser()
    parser.set_defaults()
    parser.add_option("--infile",
                      default = 'sources/sdss_gal_stack0.dat',
                      help="Input source list")
    parser.add_option("--outroot",
                      default = 'sdss_gal_stack',
                      help="Output stacked image")
    parser.add_option("--imgsize",
                      type='float',
                      default = 80,
                      help="Stack image size (GET RID OF THIS OPT)")
    parser.add_option("--imagename",
                      default = 'acis_fill_img.fits',
                      help="Counts image file name")
    parser.add_option("--expmapname",
                      default = 'expmap_fill.fits',
                      help="Exposure map file name")
    parser.add_option("--filter",
                      dest = 'filters',
                      default=[],
                      action='append',
                      help="Filter expression")
    parser.add_option('--datadir',
                      help='Data directory',
                      default='data/xdat')
    parser.add_option('--ds9',
                      default=False,
                      action='store_true',
                      help='Show images in ds9')
    (opt, args) = parser.parse_args()
    return (opt, args)

def setup_ds9():
    import ds9xpa
    xpa = ds9xpa.XPA('')
    xpa.set('tile')
    xpa.set('cmap bb')
    return xpa

def display_in_ds9(xpa):
    imgfiles = ['acis_cut_img.fits',
                'acis_fill_img.fits',
                'expmap_cut.fits',
                'expmap_fill.fits']
    xpa.set('update off')
    for iframe, imgfile in enumerate(imgfiles):
        xpa.set('frame %d' % (iframe+1))
        if os.path.exists(imgfile):
            xpa.set_cat('fits %s' % imgfile.capitalize(), filename=imgfile)
            xpa.set('zoom to fit')
            xpa.set_cat('regions', filename='ds9.reg')
        else:
            xpa.set('frame clear')

    xpa.set('update on')
    xpa.set('frame first')

def main():
    cwd = os.getcwd()
    imgtypes = ('counts', 'photflux', 'expmap', 'photflux_wexp')
    stack = dict()
    if opt.ds9:
        xpa = setup_ds9()

    import pdb
    # pdb.set_trace()

    sources = Ska.Table.read_table(opt.infile)
    filter_sources = Ska.Numpy.filter(sources, opt.filters)
    if len(filter_sources) == 0:
        print 'No matches'
        sys.exit(0)

    for source in filter_sources:
        xdat_id = '%.4f_%.4f' % (source['ra'], source['dec'])
        srcdir = join(opt.datadir,
                      xdat_id,
                      'obs%d' % source['obsid'],
                      'ccd%d' % source['ccdid'])
        os.chdir(cwd)
        if not os.path.exists(srcdir):
            continue

        os.chdir(srcdir)

        try:
            photom = pickle.load(open('apphot.pickle'))
        except IOError:
            print 'Skipping image', xdat_id, ': no apphot.dat'
            continue
        if photom[0]['NET_COUNTS'] > 15:
            print 'Skipping image', xdat_id, 'with', photom[0]['NET_COUNTS'], 'net counts'
            continue
        print 'Stacking image', xdat_id,
        image = dict()
        info = pickle.load(open('info.pickle'))

        try:
            hdus = pyfits.open(opt.imagename)
            image['counts'] = hdus[0].data.copy()
            hdus.close()

            binfactor2 = (info['ann_r1'].val * 2. / 80.)**2
            hdus = pyfits.open(opt.expmapname)
            expmap = hdus[0].data.copy()
            expmap_median = np.median(expmap.flatten())
            expmap[expmap < expmap_median * 0.25] = expmap_median
            expmap_exposure = hdus[0].header['exposure']
            image['expmap'] = expmap[0:opt.imgsize, 0:opt.imgsize] * expmap_exposure / binfactor2
            hdus.close()
            print expmap_median, expmap_exposure, binfactor2,
            if opt.ds9:
                display_in_ds9(xpa)
                raw_input()
            else:
                print
        except IOError:
            print ' File not found'
            continue

        image['photflux'] = image['counts'] / image['expmap']
        image['photflux_wexp'] = image['counts'] / image['expmap'] * expmap_exposure

        for imgtype in imgtypes:
            try:
                stack[imgtype] = stack[imgtype] + image[imgtype]
            except KeyError:
                stack[imgtype] = image[imgtype]
            except ValueError:
                print 'Size mismatch (?) for %s:' % imgtype, stack[imgtype].shape, image[imgtype].shape
                raise

    # Calc photflux_post = sum(counts) / sum(expmap)
    # [ instead of photflux = sum(counts / expmap) ]
    stack['photflux_post'] = stack['counts'] / stack['expmap']

    # Write each image out to file
    os.chdir(cwd)
    for imgtype in imgtypes + ('photflux_post',):
        pyfits.writeto(opt.outroot + '_%s.fits' % imgtype, stack[imgtype], clobber=True)


if __name__ == '__main__':
    opt, args = get_options()
    main()