#! /usr/bin/env python 

#=======================================#
#	PAW-Pipeline version: 0.6	#
#    Written by: Paul Anthony Wilson	#
# 	   paw@astro.ex.ac.uk		#
#					#
#    Pipeline for reducing CCD images	#
#=======================================#

import sys, time, shutil, subprocess
import pyraf.iraf as iraf
from os  import getcwd, chdir, remove, path, environ, getenv, makedirs, listdir, system, unlink

#####################################
# Change the parameters in this box #
#########################################################################################
#											#
#	Change to directory used to start pyraf						#
#											#
logindir = getcwd()									#
#											#
#	Your directory needs to contain the folders: bias, flat and object		#
#											#
workdir = getcwd()									#
#											#
#  	BIAS Section									#
#											#
biassection = "[11:49,1:500]"								#
trimsection = "[5:1100,5:750]"								#
#											#
#											#
#########################################################################################

from pyfits import getheader
import os.path

def cd(target):
    chdir(target)
    iraf.cd(target)

def copy_to_dir(orig, target):
    filename = path.basename(orig)
    shutil.copyfile(orig, path.join(target, filename))

def mkdir(dirname):
   "Creates a directory if it didn't exist before"
   if not os.path.isdir(dirname):
      os.mkdir(dirname)

def load_ccdred():
   iraf.imred(_doprint=0,Stdout="/dev/null")
   iraf.ccdred(_doprint=0,Stdout="/dev/null")

   iraf.unlearn("ccdproc")
   iraf.unlearn("flatcombine")
   iraf.unlearn("mkillumflat")
   iraf.unlearn("darkcombine")

   iraf.noao.imred.ccdred.ccdproc.ccdtype = ''
   iraf.noao.imred.ccdred.ccdproc.oversca = 'no'
   iraf.noao.imred.ccdred.ccdproc.trim = 'no'
   iraf.noao.imred.ccdred.ccdproc.zerocor = 'no'
   iraf.noao.imred.ccdred.ccdproc.darkcor = 'no'
   iraf.noao.imred.ccdred.ccdproc.fixpix = 'no'
   iraf.noao.imred.ccdred.ccdproc.flatcor = 'no'
   iraf.noao.imred.ccdred.ccdproc.illumcor = 'no'
   iraf.noao.imred.ccdred.ccdproc.fringec = 'no'
   iraf.noao.imred.ccdred.ccdproc.readcor = 'no'
   iraf.noao.imred.ccdred.ccdproc.scancor = 'no'
   iraf.noao.imred.ccdred.ccdproc.interac = 'no'
   iraf.noao.imred.ccdred.ccdproc.biassec = biassection
   iraf.noao.imred.ccdred.ccdproc.trimsec = trimsection

#Formatting options
bold_code = "\033[1m"
underline_code = "\033[4m"
reset_code = "\033[0;0m"

def bold(text):
   return bold_code + text + reset_code

def underline(text):
   return underline_code + text + reset_code

# Define directories
biasdir=workdir+"/bias/"
flatdir=workdir+"/flat/"
darkdir=workdir+"/dark/"
objectdir=workdir+"/object/"

#Folder names
originals = "originals"
tmp = "tmp"

# Start Pyraf
currentdir = getcwd()
# chdir(logindir)
from pyraf import iraf
# chdir(currentdir)

# Set some IRAF settings
iraf.set(imtype="fits")
load_ccdred()

print '\n'

print bold('__________________       __   ________     _________')                  
print bold('___  __ \__    |_ |     / /   ___  __ \__________  /___  ____________') 
print bold('__  /_/ /_  /| |_ | /| / /    __  /_/ /  _ \  __  /_  / / /  ___/  _ \ ')
print bold('_  ____/_  ___ |_ |/ |/ /     _  _, _//  __/ /_/ / / /_/ // /__ /  __/')
print bold('/_/     /_/  |_|___/|__/      /_/ |_| \___/\__,_/  \__,_/ \___/ \___/\n') 

print '\n Version: 0.6'

#------------------------------------------------------------------------- BIAS REDUCTION

cd(biasdir)

#Name of bias file
bias = "bias.fits"

print " "

# Check if BIAS file is present. If not, make one.
if os.path.isfile(bias):
   print "Bias file found: 	" + bold(bias)
   copy_to_dir(path.join(biasdir, bias), flatdir)
   copy_to_dir(path.join(biasdir, bias), objectdir)

else:

   print "Combining bias images...\n"
   iraf.noao.imred.ccdred.zerocombine(

###  Zerocombine parameters #####
				#
"*.fits[1]",			#		< -- You might need to change the fits cube extention
output=bias,			#
gain=1.46,			#
rdnoise=8,			#
ccdtype="",			#
combine="median",		#
reject="avsigclip",		#
process="yes",			#
delete="no",			#
scale="none")			#
				#
#################################

   iraf.imstat(bias)
   print "\n Bias file "+bold(bias)+" created\n"
   print "... done!"
   
   # Making folders if they don't already exist
   mkdir(originals)

   # Moving files to folders
   os.system( 'mv 0*.fits originals' )


# Put bias file in flats and object folder
copy_to_dir(path.join(biasdir, bias), flatdir)
#copy_to_dir(path.join(biasdir, bias), darkdir)
copy_to_dir(path.join(biasdir, bias), objectdir)


#------------------------------------------------------------------------- FLATS REDUCTION


cd(flatdir)

#Check if FLAT file is present. If not, make one.
if os.path.isfile("flat.fits"):

   print "Flat file found: 	"+bold("flat.fits")
   # Put flat file in object folder
   copy_to_dir("flat.fits", objectdir)

else:
   # Read the first .fits file and see which filter was used
   filenames = [x for x in listdir('.') if x.endswith('.fits') and x.startswith('0')]
   filter1 = getheader(filenames[0])['FILTER1']

   print "\nCombining flats taken with filter: "+underline(filter1)+"\n"
   c = [fn for fn in listdir('.') if fn.endswith('.fits') and fn.startswith('0') ]
   f = open('flats.list', 'w+')
   outf = open('bflats.list', 'w+')
   for fn in c:
      print >> f, fn + '[1,overwrite+]'
      #print >> outf, 'b' + fn
      print >> outf, fn + '[1]'
   f.close()
   outf.close()

   flats_list = '@flats.list'
   outflats_list = "@bflats.list"


   # Combine the flats
   iraf.noao.imred.ccdred.flatcombine(outflats_list,
                                      output="flat.fits",
                                      combine="median",
                                      reject="avsigclip",
                                      ccdtype="",
                                      process="yes",
                                      gain=1.46,
                                      rdnoise=8)
   iraf.imstat("flat.fits")
   print
   print " Flat file " + bold("flat.fits") +" created\n"
   print "... done!\n"

   # Making folders if they don't already exist
   mkdir(originals)
   mkdir(tmp)

   # Moving files to folders
   os.system( 'mv 0*.fits originals' )
   os.system( 'mv flats.list tmp' )
      
   # Put flat file in object folder
   copy_to_dir("flat.fits", objectdir)

#------------------------------------------------------------------------- ILLUMINATION CORRECTION


# Check if illumflat file is present. If not, make one.

if os.path.isfile("illumflat.fits"):

   print "Illumflat file found: 	"+bold("illumflat.fits")
   copy_to_dir("illumflat.fits", objectdir)

else:

   iraf.noao.imred.ccdred.mkillumflat("flat.fits",
                                      "illumflat.fits",
                                      ccdtype="",
                                      ccdproc="")
   
   # Remove bias.fits from folder
   unlink('bias.fits')
   
   # Copy illumflat.fits to object folder
   copy_to_dir("illumflat.fits", objectdir)

   iraf.imstat("illumflat.fits")
   print
   print " Illumflat file " + bold("illumflat.fits") +" created\n"
   print "... done!\n"

if os.path.isfile("blfats.list"):
   os.system( 'mv bflats.list tmp' )

'''
#------------------------------------------------------------------------- DARKS CORRECTION


cd(darkdir)

#Check if DARK file is present. If not, make one.
if os.path.isfile("dark.fits"):

   print "Dark file found: 	"+bold("dark.fits")+"\n"
   # Put flat file in object folder
   copy_to_dir("dark.fits", objectdir)

else:

   print "Combining darks..."
   c = [fn for fn in listdir('.') if fn.endswith('.fits') and fn.startswith('0') ]
   f = open('darks.list', 'w+')
   outf = open('bdarks.list', 'w+')
   for fn in c:
      print >> f, fn + '[1,overwrite+]'
      #print >> outf, 'b' + fn
      print >> outf, fn + '[1]'
   f.close()
   outf.close()

   darks_list = '@darks.list'
   outdarks_list = "@bdarks.list"


   # Combine the darks
   iraf.noao.imred.ccdred.darkcombine(outdarks_list,
                                      output="dark.fits",
                                      combine="median",
                                      reject="avsigclip",
                                      ccdtype="",
                                      process="no",
                                      gain=1.46,
                                      rdnoise=8)

   iraf.imstat("dark.fits")
   print
   print " Dark file " + bold("dark.fits") +" created\n"
   print "... done!\n"

      
   # Making folders if they don't already exist
   mkdir(originals)
   mkdir(tmp)

      # Moving files to folders
   os.system( 'mv 0*.fits originals' )
   os.system( 'mv darks.list tmp' )

   # Remove bias.fits from folder
   unlink('bias.fits')

   # Put flat file in object folder
   copy_to_dir("dark.fits", objectdir)
'''

#------------------------------------------------------------------------- REDUCE SCIENCE IMAGES

cd(objectdir)

# Check if object files are present. If not, reduce them.

if os.path.isfile("reduction-complete"):

   print "The object images have been reduced"

else:   

   print "Combining science images: "
   c = [fn for fn in listdir('.') if fn.endswith('e.fits') and fn.startswith('0') ]
   f = open('object.list', 'w+')
   outf = open('bobject.list', 'w+')
   for fn in c:
      print >> f, fn + '[1,overwrite+]'
      print >> outf, 'bfi' + fn
   f.close()
   outf.close()

   objects_list = '@object.list'
   outobjs_list = "@bobject.list"


      # Making sure no division by zero occurs
   iraf.imreplace("flat.fits",value=0.00001,upper=0,radius=0)
   iraf.noao.imred.ccdred.ccdproc(

####	OBJ parameters	     ####
				#
objects_list,			#
output=outobjs_list,		#
ccdtype="",			#
oversca="no",			#
trim="no",			#
zerocor="yes",			#
darkcor="no",			#
flatcor="yes",			#
illumco="yes",			#
fixpix="no",			#
zero=bias,			#
dark="dark.fits",		#
flat="flat.fits",		#
illum="illumflat.fits",		#
biassec=biassection,		#
trimsec=trimsection,		#
interac="no",			#
order=3)			#
				#
#################################
      
   
   # Making folders if they don't already exist
   mkdir(originals)
   mkdir(tmp)

   # Moving files to folders
   os.system( 'mv 0*.fits originals' )
   os.system( 'mv *.list tmp' )

   # Make a file which shows the reductions have been done
   os.system( 'touch reduction-complete' )

    
print bold("\nreduce.py ended successfully")+"\n"