#Python code to convert FERRE grid format spectra to individual fits format
#To be placed in the same directory as MainGrid.dat, IntermediateGrid.dat, HighGrid.dat, CaFeGrid.dat
#Will extract grids into individual .fits files with a labelling in the same order as the corresponding _Order.txt files
#ATK March 2019
#AES 27 Jan 2020 - Corrected to fits CDELT1 parameter name.

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import math as m


from numpy import loadtxt
from astropy.io import fits
from tqdm import tqdm

import os 

#get working directory

cwd=os.getcwd()

#Default number of pixels
Npix=146497

#Starting Wavelength in Angstroms
lambda1=1677.10

#Default binning of spectra (dlambda=0.05 Angstroms)
deltalambda=0.05


print 'Main Grid (3500-6000K) - Main'
print 'Intermediate T Grid (6000-8000K) - Intermediate'
print 'High T Grid (8000-10000K) - High'
print '[Ca/Fe]=0 Grid - CaFe'

#User input of which grid they would like to extract

Grid=raw_input('Enter the Grid you wish to extract (Main, Intermediate, High, CaFe or All) : ')

#Main Grid Extraction

if Grid in ['Main','main']:
	filename='MainGrid.dat'
	#make directory
	if os.path.exists(cwd+'/Main') == False:
		os.mkdir('Main')


	with open(filename,'r') as f:
		#skip the headers
		for a in range(16):
			next(f)
		#read the spectra line by line
		for index, line in enumerate(f):
			data=line.split(' ')
			data=data[0:-1]
			data=np.array(data,dtype=float)
		#write each spectrum to a .fits file
			hdu1=fits.PrimaryHDU()
			hdu1.data=data
			hdu1.header['CRPIX1']=1.0
			hdu1.header['CRVAL1']=lambda1
			hdu1.header['CDELT1']=deltalambda
			hdu1.header['NAXIS1']=Npix
			hdu1.header['CD1_1']=deltalambda
			hdu1.header['LTM1_1']=1.0   
			hdu1.writeto(cwd+'/Main/Spectrum'+str(index+1)+'.fits',overwrite=True)

#Intermediate Grid Extraction

if Grid in['Intermediate','intermediate']:
	filename='IntermediateGrid.dat'
	if os.path.exists(cwd+'/Intermediate') == False:
		os.mkdir('Intermediate')

	with open(filename,'r') as f:
		for a in range(16):
			next(f)
		for index, line in enumerate(f):
			data=line.split(' ')
			data=data[0:-1]
			data=np.array(data,dtype=float)
		

			hdu1=fits.PrimaryHDU()
			hdu1.data=data
			hdu1.header['CRPIX1']=1.0
			hdu1.header['CRVAL1']=lambda1
			hdu1.header['CDELT1']=deltalambda
			hdu1.header['NAXIS1']=Npix
			hdu1.header['CD1_1']=deltalambda
			hdu1.header['LTM1_1']=1.0   
			hdu1.writeto(cwd+'/Intermediate/Spectrum'+str(index+1)+'.fits',overwrite=True)

#High Grid Extraction

if Grid in ['High','high']:
	filename='HighGrid.dat'
	if os.path.exists(cwd+'/High') == False:
		os.mkdir('High')

	with open(filename,'r') as f:
		for a in range(16):
			next(f)

		for index, line in enumerate(f):
			data=line.split(' ')
			data=data[0:-1]
			data=np.array(data,dtype=float)
	
			hdu1=fits.PrimaryHDU()
			hdu1.data=data
			hdu1.header['CRPIX1']=1.0
			hdu1.header['CRVAL1']=lambda1
			hdu1.header['CDELT1']=deltalambda
			hdu1.header['NAXIS1']=Npix
			hdu1.header['CD1_1']=deltalambda
			hdu1.header['LTM1_1']=1.0   
			hdu1.writeto(cwd+'/High/Spectrum'+str(index+1)+'.fits',overwrite=True)


#CaFe Grid Extraction

if Grid in ['CaFe','cafe','Cafe']:
	filename='CaFeGrid.dat'
	if os.path.exists(cwd+'/CaFe')==False:
		os.mkdir('CaFe')
	
	with open(filename,'r') as f:
		for a in range(16):
			next(f)
	
		for index, line in enumerate(f):
			data=line.split(' ')
			data=data[0:-1]
			data=np.array(data,dtype=float)
	
			hdu1=fits.PrimaryHDU()
			hdu1.data=data
			hdu1.header['CRPIX1']=1.0
			hdu1.header['CRVAL1']=lambda1
			hdu1.header['CDELT1']=deltalambda
			hdu1.header['NAXIS1']=Npix
			hdu1.header['CD1_1']=deltalambda
			hdu1.header['LTM1_1']=1.0   
			hdu1.writeto(cwd+'/CaFe/Spectrum'+str(index+1)+'.fits',overwrite=True)

#If user selects all grids to be extracted

if Grid in['All','all']:
	filename=['MainGrid.dat','IntermediateGrid.dat','HighGrid.dat','CaFeGrid.dat']

	if os.path.exists(cwd+'/Main') == False:
		os.mkdir('Main')

	if os.path.exists(cwd+'/Intermediate') == False:
		os.mkdir('Intermediate')

	if os.path.exists(cwd+'/High') == False:
		os.mkdir('High')

	if os.path.exists(cwd+'/CaFe') == False:
		os.mkdir('CaFe')

	directory=['Main','Intermediate','High','CaFe']
	for i in range(0,len(filename)):
		with open(filename[i],'r') as f:
			for a in range(16):
				next(f)

			for index, line in enumerate(f):
				data=line.split(' ')
				data=data[0:-1]
				data=np.array(data,dtype=float)
	
				hdu1=fits.PrimaryHDU()
				hdu1.data=data
				hdu1.header['CRPIX1']=1.0
				hdu1.header['CRVAL1']=lambda1
				hdu1.header['CDELT1']=deltalambda
				hdu1.header['NAXIS1']=Npix
				hdu1.header['CD1_1']=deltalambda
				hdu1.header['LTM1_1']=1.0   
				hdu1.writeto(cwd+'/'+directory[i]+'/Spectrum'+str(index+1)+'.fits',overwrite=True)