#!/usr/bin/python3
import pwd, sys, os, re, shutil, time, subprocess, datetime
difx2fits_exe = "/home/lpetrov/bin/difx2fits.csh"

#
# ------------------------------------------------------------------------
#
def exe ( command ):
#"""
#    Spawn a supborcess, execute command in the context of the suborpess,
#    wait for its completion and return completion code and results.
#"""
    words = command.split()
    time_str = str(datetime.datetime.now().strftime("%Y.%m.%d_%H:%M:%S.") + "%6d" % datetime.datetime.now().microsecond).replace( " ", "0" )
    (ret, out) = subprocess.getstatusoutput ( "ulimit -s 8000000; " + command )
    return ( ret, out.split ( "\n" ) )
#
# ------------------------------------------------------------------------
#

if ( len(sys.argv) < 3 ):
     print ( "Usage: tofits difxdir fits-idi" )
     exit  ( 1 )
else:
     difxdir = sys.argv[1]
     if ( difxdir[len(difxdir):len(difxdir)] == "/" ): 
          difxdir = difxdir[0:len(difxdir)-1]
     fitsidi = sys.argv[2]

if ( not os.path.isdir ( difxdir ) ):
     print ( "Directory " + difxdir + " does not exist" )
     exit ( 1 )

for path, dirs, files in os.walk(difxdir):
    files.sort()
    for file in files:
        if ( ".calc" in file and not "#" in file ):
             file_name = path + "/" + file
             print ( "tofits Converting calc file ", file_name )
             with open(file_name,encoding="latin") as f:
                  calc_buf = f.read().splitlines()
             f.close()
             f=open(file_name,"w")
             for line in calc_buf:
                 if ( "IM FILENAME:" in line or "FLAG FILENAME:" in line):
                      finam = line.split()[2]
                      ip = finam.rindex("/")
                      if ( difxdir[-1:len(difxdir)]  == "/" ):
                            line = line.replace(finam[0:ip],difxdir[0:-1])
                      else:
                           line = line.replace(finam[0:ip],difxdir)
                 print ( line, file=f )
              
             f.close()


for path, dirs, files in os.walk(difxdir):
    files.sort()
    for file in files:
        if ( ".input" in file and not "#" in file and not ".orig" in file ):
             file_name = path + "/" + file
             file_orig_name = file_name + ".orig"
             if ( not os.path.isfile ( file_orig_name ) ):
                  os.system ( "cp " + file_name + " " + file_orig_name )
             print ( "tofits Converting input file ", file_name )
             with open(file_name,encoding="latin") as f:
                  input_buf = f.read().splitlines()
             f.close()
             f=open(file_name,"w")
             pcal_freq = 1000000
             for line in input_buf:
                 if ( line[0:20] == "PHASE CAL INT (MHZ):" ):
                      pcal_freq = int(line[20:len(line)])
                                 
             num_pcal = 0
             for line in input_buf:
                 if ( "CALC FILENAME:" in line or "OUTPUT FILENAME:" in line ):
                      finam = line.split()[2]
                      ip = finam.rindex("/")
                      if ( difxdir[-1:len(difxdir)]  == "/" ):
                            line = line.replace(finam[0:ip],difxdir[0:-1])
                      else:
                            line = line.replace(finam[0:ip],difxdir)
                      print ( line, file=f )
                 elif ( "CORE CONF FILENAME:" in line ):
                      finam = line.split()[3]
                      ip = finam.rindex("/")
                      if ( difxdir[-1:len(difxdir)]  == "/" ):
                            line = line.replace(finam[0:ip],difxdir[0:-1])
                      else:
                           line = line.replace(finam[0:ip],difxdir)
                      print ( line, file=f )
                 elif ( line[0:8] == "BW (MHZ)" ):
                      bw = float ( line.split()[3] )
                      num_pcal = int(bw/pcal_freq)
                      if ( int(bw) > num_pcal*pcal_freq ):
                           num_pcal = num_pcal + 1
#
# ------------------- In order to avoid problems at the edge of the IF
#
                      num_pcal = num_pcal - 1  
                      print ( line, file=f )
                 elif ( line[0:10] == "PHASE CALS" ):
                      if_ind = int(line.split()[2])
                      old_num_pcal = int(line.split()[4])
                      print ( "PHASE CALS %d OUT:   %d" % (if_ind, num_pcal), file=f )
                      for i in range(0,num_pcal):
                          print ( "PHASE CAL %d/%d INDEX:%d" % (if_ind, i, i), file=f )
                 elif ( line[0:10] == "PHASE CAL " and "INDEX" in line ):
                      continue
                 else:
                      print ( line, file=f )


             f.close()

tofits_com = difx2fits_exe + " -v --difx-pcal-interval 10" + " " +  difxdir + "/*.input " + fitsidi
(ret,out) = exe ( tofits_com )
if ( ret == 0 ):
     print ( "Created FITS-IDI file " + fitsidi )
     exit ( 0 )
else:
     print ( "Failure ret= ", ret, " command: ", tofits_com )
     for line in out:
         print ( line )