from __future__ import print_function
import os, sys
import numpy as np


def read_tim(timfile):
    """
    Reads a .tim file and saves non-commented lines to an array
    """
    f = open(timfile, 'r')
    lines = f.readlines()
    tim = []
    for line in lines:
        line = line.strip()
        # if not header or commented out
        if line[0] not in ['#', 'C', 'F', 'M']:
            tim.append(line)

    return np.array(tim).squeeze()

# Read .tim files
oldtim = sys.argv[1]
oldtim_array = read_tim(oldtim)
newtim = sys.argv[2]
newtim_array = read_tim(newtim)

outfile = newtim.replace('.tim', '.flagged.tim')
print("Writing to {0}".format(outfile))
with open(outfile, 'w+') as f:
    f.write("FORMAT 1 \n")
    f.write("MODE 1 \n")
    for iline in newtim_array:
       written = False
       if iline.split()[0] in ['FORMAT', 'MODE']:
           continue
       ifile = '.'.join(iline.split()[0].split('/')[-1].split('.')[0:3])
       #print(ifile)
       for line in oldtim_array:
           if written:
               continue
           lsplit = line.split()
           lsplit = np.array(lsplit).squeeze()
           if '.'.join(lsplit[0].split('.')[0:-1]) == ifile:

               # delete filename, frequency, TOA, err, obs code
               lsplit = np.delete(lsplit, [0, 1, 2, 3, 4])

               # Ignore flags already generated by pat
               ind = np.argwhere((lsplit == '-fe') +
                                 (lsplit == '-be') +
                                 (lsplit == '-f') +
                                 (lsplit == '-bw') +
                                 (lsplit == '-tobs') +
                                 (lsplit == '-tmplt') +
                                 (lsplit == '-gof') +
                                 (lsplit == '-nbin') +
                                 (lsplit == '-nch') +
                                 (lsplit == '-snr') +
                                 (lsplit == '-chan')).squeeze()

               inds = np.sort(np.concatenate((ind, ind + 1)))
               lsplit = np.delete(lsplit, inds)

               if not written:
                   # Now write to new .tim file
                   writeline = ' '.join(lsplit)
                   f.write("{0}  {1}\n".format(iline, writeline))
                   written = True

