#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import aipy
import pytz
import numpy
import pyfits
import argparse

from lsl import astro
from lsl.sim import vis as simVis
from lsl.imaging import utils, overlay
from lsl.writer.fitsidi import NumericStokes
from lsl.misc import parser as aph

import matplotlib.pyplot as plt
from matplotlib.ticker import NullFormatter


MST = pytz.timezone('US/Mountain')
UTC = pytz.UTC


NPIX_SIDE = 350


def main(args):
    filename = args.filename
    
    idi = utils.CorrelatedData(filename)
    aa = idi.getAntennaArray()
    lo = idi.getObserver()
    
    nStand = len(idi.stands)
    nChan = len(idi.freq)
    freq = idi.freq
    
    print "Raw Stand Count: %i" % nStand
    print "Final Baseline Count: %i" % (nStand*(nStand-1)/2,)
    print "Spectra Coverage: %.3f to %.3f MHz in %i channels (%.2f kHz/channel)" % (freq[0]/1e6, freq[-1]/1e6, nChan, (freq[1] - freq[0])/1e3)
    try:
        print "Polarization Products: %s" % ' '.join([NumericStokes[p] for p in idi.pols])
    except KeyError:
        # Catch for CASA MS that use a different numbering scheme
        NumericStokesMS = {1:'I', 2:'Q', 3:'U', 4:'V', 
                        9:'XX', 10:'XY', 11:'YX', 12:'YY'}
        print "Polarization Products: %s" % ' '.join([NumericStokesMS[p] for p in idi.pols])
        
    print "Reading in FITS IDI data"
    nSets = idi.integrationCount
    for set in range(1, nSets+1):
        if args.dataset != -1 and args.dataset != set:
            continue
            
        print "Set #%i of %i" % (set, nSets)
        dataDict = idi.getDataSet(set, uvMin=args.uv_min)
        
        # Build a list of unique JDs for the data
        pols = dataDict['bls'].keys()
        jdList = []
        for jd in dataDict['jd'][pols[0]]:
            if jd not in jdList:
                jdList.append(jd)
                
        # Find the LST
        lo.date = jdList[0] - astro.DJD_OFFSET
        utc = str(lo.date)
        lst = str(lo.sidereal_time())
        
        # Pull out the right channels
        toWork = numpy.where( (freq >= args.freq_start) & (freq <= args.freq_stop) )[0]
        if len(toWork) == 0:
            raise RuntimeError("Cannot find data between %.2f and %.2f MHz" % (args.freq_start/1e6, args.freq_stop/1e6))
            
        # Integration report
        print "    Date Observed: %s" % utc
        print "    LST: %s" % lst
        print "    Selected Frequencies: %.3f to %.3f MHz" % (freq[toWork[0]]/1e6, freq[toWork[-1]]/1e6)
        
        # Prune out what needs to go
        if args.include != 'all' or args.exclude != 'none':
            print "    Processing include/exclude lists"
            
            ## Create an empty output data dictionary
            newDict = {}
            for key in dataDict.keys():
                if key in ['bls', 'uvw', 'vis', 'wgt', 'msk', 'jd']:
                    newDict[key] = {}
                    for pol in dataDict[key].keys():
                        newDict[key][pol] = []
                else:
                    newDict[key] = dataDict[key]
                    
            ## Fill it
            for pol in dataDict['bls'].keys():
                for b in xrange(len(dataDict['bls'][pol])):
                    a0,a1 = dataDict['bls'][pol][b]
                    
                    if args.include != 'all':
                        if idi.stands[a0] in args.include and idi.stands[a1] in args.include:
                            for key in ['bls', 'uvw', 'vis', 'wgt', 'msk', 'jd']:
                                newDict[key][pol].append( dataDict[key][pol][b] )
                        else:
                            continue
                            
                    if args.exclude != 'none':
                        if idi.stands[a0] in args.exclude or idi.stands[a1] in args.exclude:
                            continue
                        else:
                            for key in ['bls', 'uvw', 'vis', 'wgt', 'msk', 'jd']:
                                newDict[key][pol].append( dataDict[key][pol][b] )
                                
            ## Make the substitution so that we are working with the right data now
            dataDict = newDict
            
            ## Report
            for pol in dataDict['bls'].keys():
                print "        %s now has %i baselines" % (pol, len(dataDict['bls'][pol]))
                
        # Build up the images for each polarization
        print "    Gridding"
        img1 = None
        lbl1 = 'XX'
        for p in ('xx', 'rr', 'I'):
            try:
                img1 = utils.buildGriddedImage(dataDict, MapSize=NPIX_SIDE/2, MapRes=0.5, pol=p, chan=toWork)
                lbl1 = p.upper()
            except:
                pass
                
        img2 = None
        lbl2 = 'YY'
        for p in ('yy', 'll', 'Q'):
            try:
                img2 = utils.buildGriddedImage(dataDict, MapSize=NPIX_SIDE/2, MapRes=0.5, pol=p, chan=toWork)
                lbl2 = p.upper()
            except:
                pass
                
        img3 = None
        lbl3 = 'XY'
        for p in ('xy', 'rl', 'U'):
            try:
                img3 = utils.buildGriddedImage(dataDict, MapSize=NPIX_SIDE/2, MapRes=0.5, pol=p, chan=toWork)
                lbl3 = p.upper()
            except:
                pass
                
        img4 = None
        lbl4 = 'YX'
        for p in ('yx', 'lr', 'V'):
            try:
                img4 = utils.buildGriddedImage(dataDict, MapSize=NPIX_SIDE/2, MapRes=0.5, pol=p, chan=toWork)
                lbl4 = p.upper()
            except:
                pass
                
        # Plots
        print "    Plotting"
        fig = plt.figure()
        ax1 = fig.add_subplot(2, 2, 1)
        ax2 = fig.add_subplot(2, 2, 2)
        ax3 = fig.add_subplot(2, 2, 3)
        ax4 = fig.add_subplot(2, 2, 4)
        for ax, img, pol in zip([ax1, ax2, ax3, ax4], [img1, img2, img3, img4], [lbl1, lbl2, lbl3, lbl4]):
            # Skip missing images
            if img is None:
                ax.text(0.5, 0.5, 'Not found in file', color='black', size=12, horizontalalignment='center')
                
                ax.xaxis.set_major_formatter( NullFormatter() )
                ax.yaxis.set_major_formatter( NullFormatter() )
                
                if not args.utc:
                    ax.set_title("%s @ %s LST" % (pol, lst))
                else:
                    ax.set_title("%s @ %s UTC" % (pol, utc))
                continue
                
            # Display the image and label with the polarization/LST
            cb = utils.plotGriddedImage(ax, img)
            fig.colorbar(cb, ax=ax)
            if not args.utc:
                ax.set_title("%s @ %s LST" % (pol, lst))
            else:
                ax.set_title("%s @ %s UTC" % (pol, utc))
                
            junk = img.image(center=(NPIX_SIDE/2,NPIX_SIDE/2))
            print "%s: image is %.4f to %.4f with mean %.4f" % (pol, junk.min(), junk.max(), junk.mean())
            
            # Turn off tick marks
            ax.xaxis.set_major_formatter( NullFormatter() )
            ax.yaxis.set_major_formatter( NullFormatter() )
            
            # Compute the positions of major sources and label the images
            overlay.sources(ax, aa, simVis.srcs, label=not args.no_labels)
            
            # Add in the horizon
            overlay.horizon(ax, aa)
            
            # Add lines of constant RA and dec.
            if not args.no_grid:
                if not args.topo:
                    overlay.graticleRADec(ax, aa)
                else:
                    overlay.graticleAzEl(ax, aa)
                    
        plt.show()
        
        if args.fits is not None:
            ## Loop over the images to build up the FITS file
            hdulist = [pyfits.PrimaryHDU(),]
            for img,pol in zip((img1,img2,img3,img4), (lbl1,lbl2,lbl3,lbl4)):
                if img is None:
                    continue
                    
                ### Create the HDU
                try:
                    hdu = pyfits.ImageHDU(data=img.image(center=(NPIX_SIDE/2,NPIX_SIDE/2)), name=pol)
                except AttributeError:
                    hdu = pyfits.ImageHDU(data=img, name=pol)
                    
                ### Add in the coordinate information
                hdu.header['EPOCH'] = 2000.0 + (jdList[0] - 2451545.0) / 365.25
                hdu.header['CTYPE1'] = 'RA---SIN'
                hdu.header['CRPIX1'] = NPIX_SIDE/2+1
                hdu.header['CDELT1'] = -360.0/NPIX_SIDE/numpy.pi
                hdu.header['CRVAL1'] = lo.sidereal_time()*180/numpy.pi
                hdu.header['CTYPE2'] = 'DEC--SIN'
                hdu.header['CRPIX2'] = NPIX_SIDE/2+1
                hdu.header['CDELT2'] = 360.0/NPIX_SIDE/numpy.pi
                hdu.header['CRVAL2'] = lo.lat*180/numpy.pi
                hdu.header['LONPOLE'] = 180.0
                hdu.header['LATPOLE'] = 90.0
                
                ### Add the HDU to the list
                hdulist.append(hdu)
                
            ## Save the FITS file to disk
            hdulist = pyfits.HDUList(hdulist)
            clobber = False
            if os.path.exists(args.fits):
                yn = raw_input("WARNING: '%s' exists, overwrite? [Y/n]" % args.fits)
                if yn not in ('n', 'N'):
                    clobber = True
            try:
                hdulist.writeto(args.fits, clobber=clobber)
            except IOError, e:
                print "WARNING: FITS image file not saved"
                
    print "...Done"


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description='create images from a FITS-IDI file', 
        epilog='NOTE:  If both -i/--include and -e/--exclude are specified the include list has priority.', 
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
        )
    parser.add_argument('filename', type=str, 
                        help='filename to plot')
    parser.add_argument('-1', '--freq-start', type=aph.frequency, default='10.0', 
                        help='first frequency to analyze in MHz')
    parser.add_argument('-2', '--freq-stop', type=aph.frequency, default='88.0', 
                        help='last frequency to analyze in MHz')
    parser.add_argument('-s', '--dataset', type=int, default=-1, 
                        help='data set to image')
    parser.add_argument('-m', '--uv-min', type=float, default=0.0, 
                        help='minimun baseline uvw length to include in lambda at the midpoint frequency')
    parser.add_argument('-i', '--include', type=aph.csv_int_list, default='all', 
                        help='comma seperated list of dipoles to include')
    parser.add_argument('-e', '--exclude', type=aph.csv_int_list, default='none', 
                        help='comma seperated list of dipoles to exclude')
    parser.add_argument('-t', '--topo', action='store_true', 
                        help='display an az/el grid instead of a RA/Dec grid')
    parser.add_argument('-u', '--utc', action='store_true', 
                        help='display the time as UTC, instead of LST')
    parser.add_argument('-n', '--no-labels', action='store_true', 
                        help='disable source and grid labels')
    parser.add_argument('-g', '--no-grid', action='store_true', 
                        help='disable the coordinate grid')
    parser.add_argument('-f', '--fits', type=str, 
                        help='save the images to the specified FITS image file')
    args = parser.parse_args()
    main(args)
