#!/usr/bin/env python

#
# 2014-02-07 the contour interval has been changed from 0.02 to 0.01 at L114-115
# 2014-03-10 coped with the bug in matplotlib.pyplot.clf() for mac that replot is not available 
#

import os
import sys
import pyfits
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import math
import re


def findsuffix(date, num1, num2):
    
    fits1 = '../'+date+'/rawdata/'+band+date+'_'+num1+'.fits'
    fits2 = '../'+date+'/rawdata/'+band+date+'_'+num2+'.fits'
    if os.path.exists(fits1) and os.path.exists(fits2):
        file1 = fits1
        file2 = fits2
    elif os.path.exists(fits1+'.fz') and os.path.exists(fits2+'.fz'):
        file1 = fits1+'.fz'
        file2 = fits2+'.fz'
    elif os.path.exists(fits1+'.ic') and os.path.exists(fits2+'.ic'):
        file1 = fits1+'.ic'
        file2 = fits2+'.ic'
    else:
        print 'Something wrong with compression for %s and %s' % (fits1, fits2)
        print 'ic or fz?'
        file1 = 'none'
        file2 = 'none'
 
    return file1, file2



def readimg(file1, file2):
    
    if file1.endswith('fits'):
        fim1  = pyfits.open(file1)
        im1 = fim1[0].data
        fim1.close()

        fim2 = pyfits.open(file2)
        im2 = fim2[0].data
        fim2.close()
    else:
        fim1  = pyfits.open(file1)
        im1 = fim1[1].data
        fim1.close()

        fim2 = pyfits.open(file2)
        im2 = fim2[1].data
        fim2.close()

    return im1, im2



def clipmed(img, lthresh, hthresh):

    cindex = pyfits.np.where((img > lthresh) & (img < hthresh))
    cimg = img[cindex]
    if cimg.size < 5:
        return -99

    med = pyfits.np.median(cimg)
    for i in range(5):
        sigma = math.sqrt(pyfits.np.sum((cimg - med)**2)/cimg.size)
        index = pyfits.np.where(pyfits.np.fabs(cimg - med) < 3 * sigma)
        cimg = cimg[index]
        med = pyfits.np.median(cimg)
 
    return med



def medfilter(img):

    mimg = pyfits.np.zeros((64, 64))
    # 1024x1024 -> 64x64
    for i in range(0,64):
        for j in range(0,64):
            xz=16*i
            yz=16*j
            mimg[i, j] = clipmed(img[xz:xz+16, yz:yz+16], 0, 5)
    return mimg



if __name__ == "__main__":

    argvs = sys.argv
    argc = len(argvs)
    
    if argc < 3:
        print 'Usage # twflatview.py [j|h|k]flatlist flat.fits'
        sys.exit()

    band = argvs[1][:1]
    if band != 'j' and band != 'h' and band != 'k':
        print 'The name of flat list is wrong. Not starting with j, h, or k.'
        sys.exit()
        
    fimg = pyfits.open(argvs[2])
    combdata = fimg[0].data
    fimg.close()

    # for the matplotlib bug for macosx
    mac = 0
    for item in os.uname():
        if re.search('Darwin', item):
            mac = 1
            break

    fig = plt.figure(facecolor='white', figsize=(10, 8))
    ax1 = fig.add_subplot(111)
    ax2 = fig.add_axes([0.85, 0.11, 0.03, 0.78])   #  [left, bottom, width, height]
    mycmap = cm.get_cmap('jet')
    mycmap.set_under('w')
    norm = mpl.colors.Normalize(vmin=0.94, vmax=1.06)
    cb1 = mpl.colorbar.ColorbarBase(ax2, cmap=mycmap, norm=norm, orientation='vertical', ticks=[0.94, 0.96, 0.98, 1.00, 1.02, 1.04, 1.06])
    levels = pyfits.np.arange(0.96, 1.04, 0.01)
    lwarr = pyfits.np.array([1, 1, 1, 1, 2, 1, 1, 1, 1])

    f = open(argvs[1])
    for line in f:
        v = line[:-1].split()
        date = v[0]
        num1 = v[1]
        num2 = v[2]
        file1, file2 = findsuffix(date, num1, num2)
        if file1 == 'none':
            continue

        title = band+' '+date+' '+num1+'-'+num2

        im1, im2 = readimg(file1, file2)

        im3 = im1 - im2
        med = clipmed(im3, 0, 10000)
        im3 = im3 / med
        im3 = im3 / combdata
        index = pyfits.np.where((combdata < 0.7) | (combdata > 1.5))
        im3[index] = -99

        im4 = medfilter(im3)

        ax1.imshow(im3, cmap=mycmap, vmin=0.94, vmax=1.06, interpolation='nearest', origin='lower')
        ax1.contour(im4, levels, linewidths=lwarr, colors='k', extent=(0,1024, 0,1024))
        plt.text(-13, 1.05, title)

        fig.show()

        raw_input('next:')

        # 2014-03-10 
        if mac == 1:
            plt.close()
            fig = plt.figure(facecolor='white', figsize=(10, 8))
        else:
            plt.clf()

        ax1 = fig.add_subplot(111)
        ax2 = fig.add_axes([0.85, 0.11, 0.03, 0.78])
        cb1 = mpl.colorbar.ColorbarBase(ax2, cmap=mycmap, norm=norm, orientation='vertical', ticks=[0.94, 0.96, 0.98, 1.00, 1.02, 1.04, 1.06])

    f.close()
    

