USE_MPI = 0

import numpy as np
import casacore.tables as pt
import time
import sys
import os

startfreq = 0
endfreq = 1
stepfreq = 5

#outpath = '/data/gridding/data/shortgauss_t201806301100_SBH255.binMS/'
print(sys.argv[1])
outpath = "/homes/gheller/lofar2/gheller/binMS/"+sys.argv[1]+str(startfreq)+"-"+str(endfreq)+"-"+str(stepfreq)+".binMS/"
os.mkdir(outpath)


ufile = 'ucoord.bin'
vfile = 'vcoord.bin'
wfile = 'wcoord.bin'
weights = 'weights.bin'
visrealfile = 'visibilities_real.bin'
visimgfile = 'visibilities_img.bin'
metafile = 'meta.txt'

offset = 0.0

if USE_MPI == 1:
   from mpi4py import MPI
   comm = MPI.COMM_WORLD
   rank = comm.Get_rank()
   size = comm.Get_size()
   print(rank,size)
else:
   comm = 0
   rank = 0
   size = 1

num_threads = 1

# input MS
readtime0 = time.time()

num_points_tot = 0
num_vis_tot = 0
for ifreq in range(startfreq,endfreq,stepfreq):
    freqlabel = format(ifreq,'04d')
    msfile = "/homes/gheller/lofar2/gheller/MS/"+sys.argv[1]+freqlabel+".ms/"
    ms = pt.table(msfile, readonly=True, ack=False)

    if rank == 0:
     print("Reading ", msfile)
     print("Writing ", outpath)
# load data and metadata
     with pt.table(msfile + '::SPECTRAL_WINDOW', ack=False) as freqtab:
        freq = freqtab.getcol('REF_FREQUENCY')[0] / 1000000.0
        freqpersample = np.mean(freqtab.getcol('RESOLUTION'))
        timepersample = ms.getcell('INTERVAL',0)

     print("Frequencies (MHz)   : ",freq)
     print("Time interval (sec) : ",timepersample)

     with pt.taql("SELECT ANTENNA1,ANTENNA2,sqrt(sumsqr(UVW)),GCOUNT() FROM $ms GROUPBY ANTENNA1,ANTENNA2") as BL:
        ants1, ants2 = BL.getcol('ANTENNA1'), BL.getcol('ANTENNA2')
        Ntime = BL.getcol('Col_4')[0] # number of timesteps
        Nbaselines = len(ants1)

     print("Number of timesteps : ",Ntime)
     print("Total obs time (hrs): ",timepersample*Ntime/3600)
     print("Number of baselines : ",Nbaselines)

#sp = pt.table(msfile+'::LOFAR_ANTENNA_FIELD', readonly=True, ack=False, memorytable=True).getcol('POSITION')

     ant1, ant2 = ms.getcol('ANTENNA1'), ms.getcol('ANTENNA2')

     number_of_measures = Ntime * Nbaselines
     #nm_pe_aux = int(number_of_measures / size)
     #remaining_aux = number_of_measures % size
     nm_pe = np.array(0) 
     nm_pe = int(number_of_measures / size)
     remaining = np.array(0) 
     remaining = number_of_measures % size
     print(nm_pe,remaining)

    else:
     nm_pe = None
     remaining = None

    if USE_MPI == 1:
     nm_pe = comm.bcast(nm_pe, root=0)
     remaining = comm.bcast(remaining, root=0)

# set the data domain for each MPI rank
    startrow = rank*nm_pe

    if rank == size-1:
     nm_pe = nm_pe+remaining
    print(rank,nm_pe,remaining)

    nrow = nm_pe

# read data
    uvw = ms.getcol('UVW',startrow,nrow)
    vis = ms.getcol('DATA',startrow,nrow)
    weight = ms.getcol('WEIGHT',startrow,nrow)
    print("Freqs per channel   : ",vis.shape[1])
    print("Polarizations       : ",vis.shape[2])
    print("Number of observations : ",uvw.shape[0])
    print("Data size (MB)      : ",uvw.shape[0]*vis.shape[1]*vis.shape[2]*2*4/1024.0/1024.0)
    ms.close()

# set parameters
    num_points = uvw.shape[0]
    num_points_tot = num_points_tot + num_points
    num_vis = num_points * vis.shape[1] * vis.shape[2]
    num_vis_tot = num_vis_tot + num_vis
    num_w_planes = 1 

# serialize arrays
    vis_ser_real = vis.real.flatten()
    vis_ser_img = vis.imag.flatten()
    print("data types: uvw = ",uvw.dtype," vis = ",vis_ser_real.dtype)
    #vis_ser = np.zeros(2*vis_ser_real.size)
    #for i in range(vis_ser_real.size):
    #    vis_ser[2*i]=vis_ser_real[i]
    #    vis_ser[2*i+1]=vis_ser_img[i]

    uu_ser = uvw[:,0].flatten()
    vv_ser = uvw[:,1].flatten()
    ww_ser = uvw[:,2].flatten()
    weight_ser = weight.flatten()

    print(np.amin(uu_ser),np.amax(uu_ser))
    print(np.amin(vv_ser),np.amax(vv_ser))
    print(np.amin(ww_ser),np.amax(ww_ser))

    readtime1 = time.time()

    outfile = outpath+ufile
    newFile = open(outfile, "a+b")
    binary_format = bytearray(uu_ser)
    newFile.write(binary_format)
    newFile.close()
    outfile = outpath+vfile
    newFile = open(outfile, "a+b")
    binary_format = bytearray(vv_ser)
    newFile.write(binary_format)
    newFile.close()
    outfile = outpath+wfile
    newFile = open(outfile, "a+b")
    binary_format = bytearray(ww_ser)
    newFile.write(binary_format)
    newFile.close()
    outfile = outpath+weights
    newFile = open(outfile, "a+b")
    binary_format = bytearray(weight_ser)
    newFile.write(binary_format)
    newFile.close()
    outfile = outpath+visrealfile
    newFile = open(outfile, "a+b")
    binary_format = bytearray(vis_ser_real)
    newFile.write(binary_format)
    newFile.close()
    outfile = outpath+visimgfile
    newFile = open(outfile, "a+b")
    binary_format = bytearray(vis_ser_img)
    newFile.write(binary_format)
    newFile.close()

ming = 0.0
maxg = 1.0
minw = 0.0
maxw = 1.0
outfile = outpath+metafile
f = open(outfile, 'w')
f.writelines(str(num_points_tot)+"\n")
f.writelines(str(num_vis_tot)+"\n")
f.writelines(str(vis.shape[1])+"\n")
f.writelines(str(vis.shape[2])+"\n")
f.writelines(str(Ntime)+"\n")
f.writelines(str(timepersample)+"\n")
f.writelines(str(timepersample*Ntime/3600)+"\n")
f.writelines(str(Nbaselines)+"\n")
f.writelines(str(ming)+"\n")
f.writelines(str(maxg)+"\n")
f.writelines(str(minw)+"\n")
f.writelines(str(maxw)+"\n")
f.close()


