#"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""#
#                                                          Absorption line mapping script                                       #
#"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""#

import pyfits, scipy, sys, os
from pylab import *
import math

##############################################Importing the datacube and redshift values########################################
header = raw_input('Select the datacube: ')  

while header == 'help':
    print 'A .fits datacube must be selected, for scaling reasons. Note that the datacube must be in the same folder as this script.'
    header = raw_input('Select the datacube: ')

#####Importing the redshift array#####
redshift = raw_input('Please select the redshift matrix: ')
redshift = scipy.load('%(redshift)s'%vars())

rtype = raw_input('Is the redshift matrix in pxl or kms? ')

#################Reading the datacube header##############
dcubeheader = pyfits.getheader('%(header)s'%vars())
objname = dcubeheader['OBJECT']  

axis1 = dcubeheader['NAXIS1']
axis2 = dcubeheader['NAXIS2']         #Dimension reading
axis3 = dcubeheader['NAXIS3']

print 'axis1 = ', axis1
print 'axis2 = ', axis2

startwav = dcubeheader['CRVAL3']  #reading the starting wavelength from the header, needed for spectrum adjustment
stepsize = dcubeheader['cdelt3']  #reading the stepsize, needed for conversion of redshift array to angstroms
#redshift = redshift*stepsize      #converting to angstroms

dcube = raw_input('If available, select the CR processed datacube, if not, select \'no\', the previously selected datacube will be processed.: ')

while dcube == 'help':
    print 'This program is to be used in combination with norscript.py. Which produces .npy arrays of the processed datacube to be plotted, and corresponding redshift values. These should be input as the \'processed datacube\' and \'redshift matrix\' inputs.'
    dcube = raw_input('Please select the datacube now: ')

if rtype == 'kms':
    redshift = (5000*redshift)/(3e5)
else:
    redshift = redshift*stepsize

##########################################Datacube cosmic ray processing/importing the processed numpy datacube(if available)############
if dcube == 'no':
     dcube = pyfits.getdata('%(header)s'%vars())  
     cutoff = raw_input('Datacube cosmic ray cutoff: ')
     y1 = dcube[:,0,0]

     while cutoff == 'help':
         print 'If unsure about cutoff counts, select around 800. If the final map of the galaxy shows anomalies, plot a few of the datacube spectrums where the brightness plot has the highest counts, and estimate the cutoff count from there. The attached spectrum.py can be used to plot spectrums of individual spexels.'
         cutoff = raw_input('Datacube cosmic ray cutoff: ')

     for i in scipy.arange(1,len(y1)-1,1):
         for j in range(axis2):
             for k in range(axis1):
                 if dcube[i,j,k] > cutoff:                             #setting what happens above cutoff
                     dcube[i,j,k] = cutoff
                 else:
                     if j > 0 and j < axis2 and k > 0 and k < axis1:
                         a = [dcube[i-1:i+2,j-1:j+2,k-1:k+2]]          #3x3x3 filter (for now)
                         dcube[i,j,k] = scipy.median(a)
                     elif j == 0 and k == 0:                           #top left corner
                         a = [dcube[i-1:i+2,j:j+2,k:k+2]]
                         dcube[i,j,k] = scipy.median(a)
                     elif j == (axis2-1) and k == 0:                   #bottom left corner
                         a = [dcube[i-1:i+2,j-1:j+1,k:k+2]]
                         dcube[i,j,k] = scipy.median(a)
                     elif  k == 0 and j < (axis2-1):
                         a = [dcube[i-1:i+2,j-1:j+2,k:k+2]]            #left hand side
                         dcube[i,j,k] = scipy.median(a)
                     elif j == (axis2-1) and k == (axis1-1):           #bottom right corner
                         a = [dcube[i-1:i+2,j-1:j+1,k-1:k+1]]
                         dcube[i,j,k] = scipy.median(a)
                     elif j == 0 and k == (axis1-1):                   #top right corner
                         a = [dcube[i-1:i+2,j:j+2,k-1:k+1]]
                         dcube[i,j,k] = scipy.median(a)
                     elif j == (axis2-1) and k < (axis1-1):            #bottom side
                         a = [dcube[i-1:i+2,j-1:j,k-1:k+2]]
                         dcube[i,j,k] == scipy.median(a)
                     else:
                         dcube[i,j,k] = dcube[i,j,k]
                         val = len(y1)/10
                         
                     val = len(y1)/10
                     if i%val == 0 and j == 1 and k == 1:
                         print 'Cosmic ray processing is',(i/val)*10,'% done.'
     scipy.save('%(objname)s_CRprocessedDcube'%vars(), dcube)          #saving the processed datacube for further use
     print 'Processed datacube saved as %(objname)s_CRprocessedDcube for further reference.'%vars()
else:
    dcube = scipy.load('%(dcube)s'%vars())

#################Scaling of the spectrum to the correct wavelengths##############
y = dcube[:,0,0]
yaxis = (stepsize*scipy.linspace(0, len(y)-1,num=len(y)))+startwav

####Selecting the absorption line we want to plot####
mode = raw_input('Please select mode(Hdelta_F, Hgamma_F, CN1, CN2, Ca4227, G4300, Fe4383, H_beta, Fe5015, Mg1, Mg2, Mgb, Fe5270, Fe5335, Fe5406): ')
modes = ['help', 'Hdelta_F', 'Hgamma_F', 'CN1', 'CN2', 'Ca4227', 'G4300', 'Fe4383', 'H_beta', 'Fe5015', 'Mg1', 'Mg2', 'Mgb', 'Fe5270', 'Fe5335', 'Fe5406']

while mode == help:
    print 'Select which absorption line you would like the map for.'
    mode = raw_input('Please select mode.')

while mode not in modes:
    mode = raw_input('The selected mode is invalid, please select again: ')

##############################################BANDS file######################################
#######Setting the absorption line wavelengths#####
if mode == 'Hdelta_F':
    bluelow = 4057.25
    bluehigh = 4088.5
    redlow = 4114.75
    redhigh = 4137.25
    linelow = 4091
    linehigh = 4112.25
elif mode == 'Hgamma_F':
    bluelow = 4283.5
    bluehigh = 4319.75
    redlow = 4354.75
    redhigh =4384.5
    linelow =4331.25
    linehigh =4352.25
elif mode == 'CN1':
    bluelow =4080.125
    bluehigh =4117.625
    redlow =4244.125
    redhigh =4284.125
    linelow =4142.125
    linehigh =4177.125
elif mode == 'CN2':
    bluelow =4083.875
    bluehigh =4096.375
    redlow =4244.125
    redhigh =4284.125
    linelow =4142.125
    linehigh =4177.125
elif mode == 'Ca4227':
    bluelow =4211
    bluehigh =4219.75
    redlow =4241
    redhigh =4251
    linelow =4222.25
    linehigh =4235.75
elif mode == 'G4300':
    bluelow =4266.375
    bluehigh =4282.625
    redlow =4318.875
    redhigh =4335.125
    linelow =4281.375
    linehigh =4316.375
elif mode == 'Fe4383':
    bluelow =4359.125
    bluehigh =4370.375
    redlow =4442.875
    redhigh =4455.375
    linelow =4369.125
    linehigh =4420.375
elif mode == 'H_beta':
    bluelow =4827.875
    bluehigh =4847.875
    redlow =4876.625
    redhigh =4891.625
    linelow =4847.875
    linehigh =4876.625
elif mode == 'Fe5015':
    bluelow =4946.5
    bluehigh =4977.750
    redlow =5054
    redhigh =5065.250
    linelow =4977.750
    linehigh =5054
elif mode == 'Mg1':
    bluelow =4895.125
    bluehigh =4957.625
    redlow =5301.125
    redhigh =5366.125
    linelow =5069.125
    linehigh =5134.125
elif mode == 'Mg2':
    bluelow =4895.125
    bluehigh =4957.625
    redlow =5301.125
    redhigh =5366.125
    linelow =5154.125
    linehigh =5196.625
elif mode == 'Mgb':
    bluelow =5142.625
    bluehigh =5161.375
    redlow =5191.375
    redhigh =5206.375
    linelow =5160.125
    linehigh =5192.625
elif mode == 'Fe5270':
    bluelow =5233.150
    bluehigh =5248.150
    redlow =5285.65
    redhigh =5318.15
    linelow =5245.65
    linehigh =5285.65
elif mode == 'Fe5335':
    bluelow =5304.625
    bluehigh =5315.875
    redlow =5353.375
    redhigh =5363.375
    linelow =5312.125
    linehigh =5352.125
elif mode == 'Fe5406':
    bluelow =5376.25
    bluehigh =5387.5
    redlow =5415
    redhigh =5425
    linelow =5387.5
    linehigh =5415

matrix = scipy.zeros((axis2,axis1))

for i in range(axis2):
    for j in range(axis1):
        if redshift[i][j] > 300:
            redshift[i][j] = 50     #needs to be improved
            
        y = dcube[:,i,j]

        blow = (bluelow-startwav+redshift[i][j])/stepsize
        bhigh = (bluehigh-startwav+redshift[i][j])/stepsize
        rlow = (redlow-startwav+redshift[i][j])/stepsize
        rhigh = (redhigh-startwav+redshift[i][j])/stepsize
        llow = (linelow-startwav+redshift[i][j])/stepsize
        lhigh = (linehigh-startwav+redshift[i][j])/stepsize

        #if blow < 0 and bhigh < 0:
        #    print 'The limits are below the starting wavelength.'
        #    print blow, bhigh
            #sys.exit(1)

        #if blow > bhigh or blow < 0 or bhigh < 0:
        #    blow = 0
        #    bhigh = 10
       
        #if rlow > rhigh or rlow < 0 or rhigh < 0:
        #    rlow = 10
        #    rhigh = 20
        
        #if llow > lhigh or llow < 0 or lhigh < 0:
        #    llow = 20
        #    lhigh = 30

        blue = scipy.mean(y[blow:bhigh])
        red = scipy.mean(y[rlow:rhigh])
        line = scipy.mean(y[llow:lhigh])

        #########Linear interpolation between the red and blue sides, defining them as points at their mean values##########
        #blue side coordinates

        xcoordblue = (blow+bhigh)/2 
        ycoordblue = blue
        #red side coordinates
        xcoordred = (rlow+rhigh)/2
        ycoordred = red     
        #absorption line x coordinate(line-width mean)
        xcoordline = (llow+lhigh)/2
        
        xdistance = xcoordred - xcoordblue
        ydistance = ycoordred - ycoordblue
        gradient = ydistance/xdistance
        
        ycoordline = ycoordblue + gradient*(xcoordline-xcoordblue)
        
        depth = line/ycoordline
                
        matrix[i][j] = depth

        if matrix[i][j] > 1:
            matrix[i][j] = 1

######Finding/creating the folder for the array and plot to be saved in
def ensure_dir(f):
    d = os.path.dirname(f)
    if not os.path.exists(d):
        os.mkdir(d)

ensure_dir('%(objname)s/'%vars())
os.chdir('%(objname)s/'%vars())
ensure_dir('abswidth/'%vars())
os.chdir('abswidth/'%vars())

scipy.save('%(objname)s_%(mode)s'%vars(),matrix)

#matrix histogram for upper and lower limit selection
title('Histogram for selecting the upper and lower limits of the final map.')
hist(matrix, bins = 40)
show()

vMin = input('Please select lower limit: ')
vMax = input('Please select upper limit: ')

while vMin > vMax:
    print 'The lower limit cannot be higher that the upper limit, please select again.'
    vMin = input('Please re-select lower limit: ')
    vMax = input('Please re-select upper limit: ')

print 'Datacube shape: ', dcube.shape

title('%(objname)s %(mode)s absorption depth'%vars())
pcolormesh(matrix, vmin = vMin, vmax = vMax)
colorbar()
savefig('%(objname)s_%(mode)s_absdepth.png'%vars())
show()