# -*- coding: utf-8 -*-
"""
    Created on Nov 25 2016
    @author: Eder Martioli
    Laboratorio Nacional de Astrofisica, Brazil
    
    spectrum.py is a library of classes and functions to
    handle spectral data.
    """

import numpy as np
from scipy import constants

########## SPECTRUM CLASS ############
class Spectrum :
    'Common base class for a spectrum'
    
    def __init__(self, Filename):

        """
        Create a Spectrum object.
        
        Parameters
        ----------
        filename : string
            File to read the spectrum from.

        Examples
        --------

        >>> sp = Spectrum("spectrumfile.1d.spc")
        """
        self.filename = Filename
        self.load_spectrum(self.filename)

    def load_spectrum(self,Filename):
        try:
            self.wl,self.flux,self.var = np.loadtxt(Filename, unpack=True, comments='#', usecols=(0,1,2), delimiter=' ')
        except:
            print "Error: could not open file:",Filename
            exit()

    def getdata(self, wl0=0., wlf=0.) :
        """
            Retrieve data for a given wavelength range
            
            Parameters
            ----------
            wl0 : initial wavelength       [nm]
            wlf : final wavelength         [nm]
            
            Return : wl[], flux[], variance[]
            """
        
        if (wl0 == 0.) :
            wl0 = self.wl[0]
        if (wlf == 0.) :
            wlf = self.wl[-1]

        mask = np.where((self.wl > wl0) & (self.wl < wlf))

        return self.wl[mask],self.flux[mask],self.var[mask]

    def applyRVShift(self, RVshift, interp=False) :
        """
            Apply radial velocity shift to the wavelength data.
            
            Parameters
            ----------
            RVshift : radial velocity shift      [m/s]
            interp  : interpolate shifted data to keep original wavelength sampling? [boolean]
            """
        self.rvshit = RVshift
        
        if interp == True :
            wl_tmp = self.wl*(1.0 + self.rvshit/constants.c)
            flux_tmp = np.interp(self.wl, wl_tmp, self.flux)
            self.flux = flux_tmp
        else :
            self.wl *= (1.0 + self.rvshit/constants.c)


