import subprocess
import sys
import time

# input data
datadir = "cgheller@login.m100.cineca.it://m100_scratch/userexternal/cgheller/Lofar/data_model"
cleandir = "cgheller@login.m100.cineca.it://m100_scratch/userexternal/cgheller/Lofar/data_wsclean"
dirtydir = "cgheller@login.m100.cineca.it://m100_scratch/userexternal/cgheller/Lofar/data_wsdirty"
filename_prefix = "RADIO30MHz"
file_ext = ".fits"

# measurement set
workdir = "."
ms_name = "hba-10hrs_t201806301100_SBH255.MS"    #"visibilities_aux.MS"
ms_tar = workdir+"/"+ms_name+"-template.tar"

# general parameters
resolution = "2.0asec"
size = "2000"

# command for creating visibilities
vis_command = "wsclean"

# parameters for the noise
noise_command = "../useful_scripts/noise.py"
noise_factor = "0.0001"

# command for imaging
img_command = "wsclean"
img_parallel = "4"
img_iter = "10000"
#img_options = "-weight briggs 0 -taper-gaussian 60 -apply-primary-beam -reorder -niter 10000 -mgain 0.8 -auto-threshold 5"
#img_options = "-apply-primary-beam -reorder -niter 10000 -mgain 0.8 -auto-threshold 5"

## create empty MS
## LBA 30 MHz
#/opt/losito/bin/synthms --name ms_name --tobs 1 --ra 1.0 --dec 1.57 --lofarversion 1 --station LBA -minfreq 30e6 --maxfreq 30e6
## HBA 150 MHz
#/opt/losito/bin/synthms --name dataset --tobs 1 --ra 1.0 --dec 1.57 --lofarversion 1 --station HBA --minfreq 150e6 --maxfreq 150e6
#
#tar cvf ms_tar ms_name

tinverse = 0.0
tcopy = 0.0
tnoise = 0.0
timaging = 0.0
tget = 0.0
tput = 0.0

for id in range(1,2):
    imgname = filename_prefix+str(id)
    #fitsfilename = datadir+'/'+imgname+file_ext
    fitsfilename = datadir+'/'+filename_prefix+file_ext
    commandline = "scp "+datadir+"/"+imgname+file_ext+" ."
    t0 = time.time()
    error = subprocess.call(commandline,shell=True)
    error = subprocess.run(["cp",imgname+file_ext,"test-model.fits"])
    tget = tget+(time.time()-t0)

    # untar the template MS
    error = subprocess.run(["tar","xvf",ms_tar])

    # create the visibilities
    t0 = time.time()
    error = subprocess.run([vis_command,"-j",img_parallel,"-scale",resolution,"--predict","-use-idg","-grid-with-beam","-name", "test", ms_name])
    tinverse = tinverse+(time.time()-t0)

    # Copy the MODEL_DATA column int the DATA column
    t0 = time.time()
    error = subprocess.run(["python","../useful_scripts/copycolumn.py",ms_name])
    tcopy = tcopy+(time.time()-t0)

    # Add noise
    t0 = time.time()
    error = subprocess.run([noise_command,"--msin",ms_name,"--factor",noise_factor])
    tnoise = tnoise+(time.time()-t0)

    # Imaging
    t0 = time.time()
    error = subprocess.run([img_command,"-j",img_parallel,"-apply-primary-beam","-reorder","-niter",img_iter,\
                            "-mgain", "0.8", "-auto-threshold", "5", "-size",size,size,"-scale",resolution,\
                            "-name",imgname,ms_name])
    timaging = timaging+(time.time()-t0)

    # remove at the end the measurement set
    error = subprocess.run(["rm","-fr",workdir+"/"+ms_name])
    error = subprocess.run(["rm","test-model.fits"])
    error = subprocess.run(["mkdir",imgname])
    #commandline = "rm *beam*.fits"
    #error = subprocess.Popen(commandline, shell=True,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
    t0 = time.time()
    commandline = "scp "+imgname+"-dirty.fits"+" "+dirtydir
    error = subprocess.call(commandline, shell=True)
    commandline = "scp "+imgname+"-image-pb.fits"+" "+cleandir
    error = subprocess.call(commandline, shell=True)
    tput = tput+(time.time()-t0)
    commandline = "mv *.fits "+imgname
    error = subprocess.call(commandline, shell=True)
##    error = subprocess.run(["tar","cvf",imgname+".tar",imgname])
    error = subprocess.run(["rm","-fr",imgname])
    #commandline = "tar cvf "+imgname+".tar *.fits"
    #error = subprocess.Popen(commandline, shell=True,stdout=subprocess.PIPE,stderr=subprocess.PIPE)
    #commandline = "rm *.fits"
    #error = subprocess.Popen(commandline, shell=True,stdout=subprocess.PIPE,stderr=subprocess.PIPE)


print("Summary of timings: ")
print("Get data (sec)     = ",tget)
print("Put data (sec)     = ",tput)
print("Inverse time (sec) = ",tinverse)
print("Copy time (sec)    = ",tcopy)
print("Noise time (sec)   = ",tnoise)
print("Imaging time (sec) = ",timaging)