import math
import numpy as np
from scipy import signal
from astropy.io import fits
import  tarfile
from subprocess import call
import os
from scipy.interpolate import interp1d
import numpy as np
from astropy.table import Table
import sys
import argparse

parser = argparse.ArgumentParser(description='Converter for carbon stars models files')

parser.add_argument("input_file", help="input tar file (without .tar !!!)")
parser.add_argument("fwhm", help="FWHM of convolution kernel")
parser.add_argument("output_dir", help="path to directory for output fits files")
args = parser.parse_args()


file_names = []
name_tar = args.input_file
tar = tarfile.open(name_tar+".tar")
for member in tar.getmembers():
    f_ = tar.extractfile(member)
    if '.spe' in f_.name:
      file_names.append(f_.name)
tar.close()

#file_name = 'mxcom10v06_t2500_g+000_m0100_o88029_c88241_fe75630_xi25_c11a.spe'
for file_name in file_names:
  
  i = 0
  flux = []
  gauss_core = []

  length_convolve=101
  med = length_convolve/2
  sigma = float(args.fwhm) /2.355
  K = (1/(sigma*math.sqrt(2*math.pi)))
  for i in range(length_convolve):
    val = K * math.exp(-(i-med)**2/(2*sigma**2))
    gauss_core.append(val)

  gauss_core = np.array(gauss_core)
  gauss_core = gauss_core/sum(gauss_core)


  wavelengths=[]
  data_array=[]
  os.system("tar -xvf "+name_tar+".tar "+file_name)
  f = open(file_name)
  for line in f:
    spltdstr=line.split( )
    data_array.append([float(spltdstr[0]),float(spltdstr[2])])
  f.close()
  os.system("rm "+file_name)
  data_array = np.array(data_array)
  data_array = data_array[data_array[:,0].argsort()]
  
  flux = data_array[: , 1]
  wavelengths = data_array[: , 0]
  convolved = np.convolve(flux, gauss_core,'same')
  wavelengths = np.array(wavelengths)
  wavelengths_atm = wavelengths / (1.0 + 2.735182e-4 + 131.4182/wavelengths**2 + 2.76249e8/wavelengths**4)
  print max(wavelengths),min(np.log10(wavelengths_atm)),min(wavelengths[:-60])
  intr_pol = interp1d(np.log10(wavelengths_atm), convolved, kind='linear')
  fluxes_atm = intr_pol(np.log10(wavelengths[:-60]))
  #convolved = flux

  wl_last = 0

  '''wl=np.zeros(int(len(convolved)//2))
  binned=np.zeros(int(len(convolved)//2))
  for j in range(len(convolved)):
    n = j//2
    binned[n]+=convolved[j]/2
    wl[n]+=wavelengths[2*j//2]/2'''
  binned = convolved
  wl = wavelengths
    
    
  fits_name_vac = file_name.replace(".spe", "_vac.fits")
  fits_name_atm = file_name.replace(".spe", "_atm.fits")

  os.chdir(args.output_dir+'/'+name_tar)


  dat_output = np.zeros((len(binned),2))
  dat_output[:,0] = binned
  dat_output[:,1][:-60] = fluxes_atm
  #print fname
  fits.HDUList([fits.PrimaryHDU(data=binned,header=None)]).writeto(fits_name_vac, clobber=True)
  fits.HDUList([fits.PrimaryHDU(data=fluxes_atm,header=None)]).writeto(fits_name_atm, clobber=True)




  name = file_name
  i = name.find('_t')+2
  temp = float(name[i:(i+4)])
  #print temp
  fits.setval(fits_name_vac, 'CASTEFF', value=temp)
  fits.setval(fits_name_atm, 'CASTEFF', value=temp)

  #LOG(g)
  i = name.find('_g')+2
  logg = float(name[i:(i+3)])
  logg = logg/100.
  #print logg
  fits.setval(fits_name_vac, 'CASLOGG', value=logg)
  fits.setval(fits_name_atm, 'CASLOGG', value=logg)
  #MASS of star
  i = name.find('_m')+2
  mass = float(name[i:(i+4)])
  mass = mass /100.
  #print mass
  fits.setval(fits_name_vac, 'CASMASS', value=mass)
  fits.setval(fits_name_atm, 'CASMASS', value=mass)

  #Oxygen abundance
  i = name.find('_o')+2
  abund_o = float(name[i:(i+5)])
  abund_o = abund_o /1e4
  #print abund_o
  fits.setval(fits_name_vac, 'CASAB_O', value=abund_o)
  fits.setval(fits_name_atm, 'CASAB_O', value=abund_o)

  #Carbon abundance
  i = name.find('_c')+2
  abund_c = float(name[i:(i+5)])
  abund_c = abund_c /1e4
  #print abund_c
  fits.setval(fits_name_vac, 'CASAB_C', value=abund_c)
  fits.setval(fits_name_atm, 'CASAB_C', value=abund_c)

  #Iron abundance
  i = name.find('_fe')+3
  abund_fe = float(name[i:(i+5)])
  abund_fe = abund_fe /1e4
  #print abund_fe
  fits.setval(fits_name_vac, 'CASAB_FE', value=abund_fe)
  fits.setval(fits_name_atm, 'CASAB_FE', value=abund_fe)

  #Microturbulent velocity
  i = name.find('_xi')+3
  mic_turb = float(name[i:(i+1)])
  mic_turb = mic_turb /10.
  #print mic_turb
  fits.setval(fits_name_vac, 'CASXI_M', value=mic_turb)
  fits.setval(fits_name_atm, 'CASXI_M', value=mic_turb)

  fits.setval(fits_name_vac, 'ORIG', value=file_name)
  fits.setval(fits_name_atm, 'ORIG', value=file_name)


  CovO = 10.**(abund_c-abund_o)
  fits.setval(fits_name_vac, 'CASC_O', value=CovO)
  fits.setval(fits_name_atm, 'CASC_O', value=CovO)

  FeSurH = (abund_fe-7.563)
  fits.setval(fits_name_vac, 'CASFE_H', value=FeSurH)
  fits.setval(fits_name_atm, 'CASFE_H', value=FeSurH)

  CmO = -1.e-10; 
  if abund_c > abund_o:
    CmO = math.log10(10.**(abund_c) - 10.**(abund_o))
  fits.setval(fits_name_vac, 'CASCEX', value=CmO)
  fits.setval(fits_name_atm, 'CASCEX', value=CmO)

  fits.setval(fits_name_vac, 'CRPIX1', value=1)
  fits.setval(fits_name_atm, 'CRPIX1', value=1)

  fits.setval(fits_name_vac, 'CTYPE1', value='WAVE-LOG')
  fits.setval(fits_name_atm, 'CTYPE1', value='WAVE-LOG')

  fits.setval(fits_name_vac, 'CDELT1', value=1e-5)
  fits.setval(fits_name_atm, 'CDELT1', value=1e-5)

  fits.setval(fits_name_vac, 'CRVAL1', value=math.log(wl[0]))
  fits.setval(fits_name_atm, 'CRVAL1', value=math.log(wl[0]))

  os.chdir('../../')