#!/bin/python
'''
Author: Yin Lin        Date: September 23, 2013

The class object StarClass loads stars from the catalog.py and from the source extractor cat file and converts them into proper format.
The lists of stars are then passed to the subclass, StarCalibration, to cross match two lists and perform offset and rotational calibration 
by changing the header keywords of the corresponding arguments. In order to use the code, you only have to call StarCalibration which inherits 
from StarClass.

'''

from functions import *
import pyfits
import matplotlib.pyplot as plt
import numpy as np
from astropy import wcs
from scipy.optimize import leastsq,fsolve
import warnings
import os
import subprocess
from glob import *
import time
from ds9 import *
import PyGuide as pg

#ignore the warning caused by astropy
warnings.filterwarnings("ignore")
    
class StarClass(object):
    """Load a list of star in the fits image and initilize stars from catalog if specified"""

    def __init__(self,fitsImageName,fitsTableName=None,manual=False,fitsCatalogName=None,caldir='./cal/',fdir='./origin/',sedir='./config/',manCat=None,manCatFile=None):
        """
        Keyword arguments:
        fitsImageName -- the input fits image file
        fitsTableName -- the input fits table file from catalog.py. Note that it has to be in the default format(asu-fits)
        manual -- if False, the program uses Sextractor to extract stars in the image, else, the program uses ds9 and PyGuide 
                  to manually calibrate the image. Also, if True, fitsCatalogName must be specified.
        fitsCatalogName -- t
he catalog image from catalog for manual calibration
        """

        self.fitsImageName = fitsImageName
        self.fitsTableName = fitsTableName
        self.fitsCatalogName = fitsCatalogName
        self.catalog = []
        self.starList = []
        self.sedir = sedir
        self.fdir = fdir
        self.caldir = caldir
        self.fitsdir = self.fdir + self.fitsImageName
        self.manual = manual
        self.minError = 2000
        self.calibrate = True
        self.manCat = ma
nCat
        self.manCatFile = manCatFile
        
        #test to see if calibration is necessary
        imageList = pyfits.open(self.fitsdir)
        header = imageList[0].header
        
        try:
            #if calibrated file with suffix _offCal_rotCal already exists, skip the calibration 
            imageListCal = pyfits.open(self.caldir+self.fitsImageName[:-5]+'_offCal_rotCal.fits')
            headerCal = imageListCal[0].header
            CALERR = headerCal['CALERR']
            print 'CALERR:' + str(CALERR)  
            if CALERR > self.minError:
                raise
            self.calibrate = False
        except:
            pass
        
        if self.manual and self.calibrate :
            #defining basic parameters for PyGuide.findStars. More info can be found on help(pyGuide.findStars)
            ccd = pg.CCDInfo(0,0.00001,1,2500)

            satMask = np.zeros(image.shape)
            mask = np.zeros(image.shape)

            #this returns Centroid class instance in PyGuide
            centroidData = pg.findStars(image,mask=mask,satMask=satMask,rad=30,ccdInfo=ccd,doDS9=False)[0]

            #sort stars and discard the one that does not make sense, e.g, the stars that are out of array
            starList = []
            _count = 0
            for data in range(len(centroidData)):    
                if centroidData[_count].xyCtr[0] > 0 and centroidData[_count].xyCtr[1] > 0:
                    starList.append(centroidData[_count].xyCtr)
                _count += 1    
            
        if not self.manual and self.calibrate:
            
            #use source extractor to extractor sources, visit man page for more info
            catName = self.caldir + self.fitsImageName[:-5] + '.cat'
            paramdir = self.sedir + 'default.sex'
            checkimg = self.caldir + self.fitsImageName[:-5] + '.check'
            proc = subprocess.Popen(['sex',self.fdir+fitsImageName,'-c',paramdir,'-CATALOG_NAME',catName,'-CHECKIMAGE_NAME',checkimg])
            proc.communicate()
            #read cat file to extract stars
            starList = readCat(catName)
        
        #coordinates in pixels
        if self.calibrate:
            self.starList = np.array(starList)
    
        #if fits table is provided, loading fits table and convert the list of star into standard numpy array
        if self.manCat == 'full':
            catalog = []
            catalogAppend = readCat(self.manCatFile)
            for star in catalogeAppend:
                catalog.append([star[0],star[1]])
        
        elif self.fitsTableName != None:
            tableList = pyfits.open(self.fitsTableName)
            
            h1 = tableList[1]
            catalog = h1.data
            catalog_x = []
            catalog_y = []
            for i in range(len(catalog)):
                catalog_x.append(catalog[i][0])
                catalog_y.append(catalog[i][1])
            
            #if semi manual catalog is used, add stars in manCatFile into the existing catalog
            if self.manCat == 'semi':
                for star in readCat(self.manCatFile):
                    catalog_x.append(star[0])
                    catalog_y.append(star[1])
            catalog = zip(catalog_x,catalog_y)
        
        self.catalog = np.array(catalog)
        print len(self.catalog)
                
        #always remember to close the file to avoid memory leakage
        imageList.close()
        
    '''
    def showCatalogPlot(self):
        #show a plot of stars in catalog
        
        if self.fitsTableName == None:
        
            print 'No Fits Table Provided!'
            
        else:       
            x = []
            y = []
            for i in range(len(self.catalog)):
                #print catalog[i][0],catalog[i][1]
                x.append(self.catalog[i][0])
                y.append(self.catalog[i][1])
            plt.plot(x,y,'ro')
            plt.show()      
            
    def showImagePlot(self):
        #show a plog of stars in fits image found by the pyGuide

        starx = []
        stary = []
        for i in range(len(self.starList)):
            starx.append(self.starList[i][0])
            stary.append(self.starList[i][1])
        plt.plot(starx,stary,'ro')
        plt.gca().invert_xaxis()    
        plt.show()
    '''

class StarCalibration(StarClass):
    #Calibration utility for fits image. Catalog data needs to be provided

    def __init__(self,fitsImageName,fitsTableName,fitsCatalogName=None,manual=False,paramFile=None,caldir='./cal/',fdir='./origin/',sedir='./config/',height=3,manCat=None,manCatFile=None):
        
        #Initialize class attributes
        self.paramFile = paramFile 
        self.calibratedStar = []
        self.fitsImageName = fitsImageName
        self.fitsTableName = fitsTableName
        self.fitsCatalogName = fitsCatalogName
        self.calpix = []
        self.pix = []
        self.labelOn = True
        self.sedir = sedir
        self.caldir = caldir
        self.fdir = fdir
        self.fitsdir = self.fdir + self.fitsImageName
        self.height = height
        self.manual = manual
        
		#create folder for calibrated files if doesnt exit
        if not os.path.exists(self.caldir):
            os.makedirs(self.caldir)		
		
		#ignore the warning caused by astropy
        warnings.filterwarnings("ignore")
        
        #Initialize the super class, StarClass
        super(StarCalibration,self).__init__(self.fitsImageName,self.fitsTableName,self.manual,self.fitsCatalogName,self.caldir,self.fdir,self.sedir,manCat=manCat,manCatFile=manCatFile)
        
        #calibration needs at least 2 star to perform. Raise error if it has fewer than 2 star detected. This becomes a warning in manage.py
        if len(self.starList) < 2 and self.calibrate:        
            raise ValueError, '2 of more stars required to calculate calibration parameters. Only %s star/s detected!' %len(self.starList)
        
        #Initialize reference pixels from the header
        imageL