from scipy.interpolate import InterpolatedUnivariateSpline as interp

import os

import sys

import numpy as np

import DataStructures

import matplotlib.pyplot as plt

import Correlate

import FitsUtils

import FindContinuum

# import Units

from astropy import units, constants

import FittingUtilities

homedir = os.environ["HOME"]

modeldir = homedir + "/School/Research/Models/Sorted/Stellar/Vband/"

#Define regions contaminated by telluric residuals or other defects. We will not use those regions in the cross-correlation

badregions = [[588.8, 589.9],

[627.1, 635.4]]

badregions = [[0, 466],

[759, 9e9]]

#Set up model list

model_list = [modeldir + "lte30-4.00-0.0.AGS.Cond.PHOENIX-ACES-2009.HighRes.7.sorted",

modeldir + "lte78-4.0+0.5.Cond.PHOENIX2004.direct.7.sorted"]

star_list = []

temp_list = []

gravity_list = []

metal_list = []

model_data = []

for fname in model_list:

if "PHOENIX2004" in fname:

temp = int(fname.split("lte")[-1][:2]) * 100

gravity = float(fname.split("lte")[-1][3:6])

metallicity = float(fname.split("lte")[-1][6:10])

elif "PHOENIX-ACES" in fname:

temp = int(fname.split("lte")[-1][:2]) * 100

gravity = float(fname.split("lte")[-1][3:7])

metallicity = float(fname.split("lte")[-1][7:11])

print "Reading in file %s" % fname

x, y = np.loadtxt(fname, usecols=(0, 1), unpack=True)

model_data.append(DataStructures.xypoint(x=x * units.angstrom.to(units.nm) / 1.00026, y=10 ** y))

star_list.append(str(temp))

temp_list.append(temp)

gravity_list.append(gravity)

metal_list.append(metallicity)

if __name__ == "__main__":

#Parse command line arguments:

fileList = []

extensions = True

tellurics = False

trimsize = 100

smooth = False

for arg in sys.argv[1:]:

if "-e" in arg:

extensions = False

if "-t" in arg:

tellurics = True #telluric lines modeled but not removed

if "-s" in arg:

smooth = True

else:

fileList.append(arg)

for fname in fileList:

if extensions:

orders = FitsUtils.MakeXYpoints(fname, extensions=extensions, x="wavelength", y="flux", errors="error")

if tellurics:

model_orders = FitsUtils.MakeXYpoints(fname, extensions=extensions, x="wavelength", y="model")

for i, order in enumerate(orders):

orders[i].cont = FindContinuum.Continuum(order.x, order.y, lowreject=2, highreject=2)

orders[i].y /= model_orders[i].y

else:

orders = FitsUtils.MakeXYpoints(fname, errors=2)

numorders = len(orders)

for i, order in enumerate(orders[::-1]):

#Only use the middle half of each order (lots of noise on the edges)

DATA = interp(order.x, order.y)

CONT = interp(order.x, order.cont)

ERROR = interp(order.x, order.err)

left = int(order.size() / 4.0)

right = int(order.size() * 3.0 / 4.0 + 0.5)

order.x = np.linspace(order.x[left], order.x[right], right - left + 1)

order.y = DATA(order.x)

order.cont = CONT(order.x)

order.err = ERROR(order.x)

#Remove bad regions from the data

for region in badregions:

left = np.searchsorted(order.x, region[0])

right = np.searchsorted(order.x, region[1])

if left == 0 or right == order.size():

order.x = np.delete(order.x, np.arange(left, right))

order.y = np.delete(order.y, np.arange(left, right))

order.cont = np.delete(order.cont, np.arange(left, right))

order.err = np.delete(order.err, np.arange(left, right))

else:

print "Warning! Bad region covers the middle of order %i" % i

print "Interpolating rather than removing"

order.y[left:right] = order.cont[left:right]

order.err[left:right] = 9e9

#Remove whole order if it is too small

remove = False

if order.x.size <= 1:

remove = True

else:

velrange = 3e5 * (np.median(order.x) - order.x[0]) / np.median(order.x)

if velrange <= 1050.0:

remove = True

if remove:

print "Removing order %i" % (numorders - 1 - i)

orders.pop(numorders - 1 - i)

else:

order.cont = FittingUtilities.Continuum(order.x, order.y, lowreject=3, highreject=3)

orders[numorders - 1 - i] = order.copy()

"""

for i, order in enumerate(orders):

plt.plot(order.x, order.y/order.cont+i)

plt.show()

sys.exit()

"""

#Smooth data

if smooth:

for order in orders:

order.y /= FittingUtilities.savitzky_golay(order.y, 91, 5)

output_dir = "Cross_correlations/"

outfilebase = fname.split(".fits")[0]

if "/" in fname:

dirs = fname.split("/")

output_dir = ""

outfilebase = dirs[-1].split(".fits")[0]

for directory in dirs[:-1]:

output_dir = output_dir + directory + "/"

output_dir = output_dir + "Cross_correlations/"

#Do the cross-correlation

for vsini in [10, 20, 30, 40]:

Correlate.PyCorr2(orders, resolution=60000, outdir=output_dir, models=model_data, stars=star_list,

temps=temp_list, gravities=gravity_list, metallicities=metal_list,

vsini=vsini * units.km.to(units.cm), debug=False, outfilebase=outfilebase)