import os,psutil,time,datetime as dttime
from optparse import OptionParser
from paircarstools.access_ms import *
from paircarstools.basic_func import *
from paircarstools.flagger import *
from paircarstasks.split_ms import *
from casatasks import *
from casatools import msmetadata,table,measures,quanta,agentflagger,image,calibrater,ms
import multiprocessing as mp
from multiprocessing import current_process
os.system('rm -rf casa*log')

'''
Code is written by Devojyoti Kansabanik , 20 May, 2021
'''

def split_ms_parallel(paramList,cpu_frac):
	'''
	Function to run multiple split jobs in parallel from same or different measurement sets
	Parameters
	----------
	paramList : list 
		Lists of list of parameters for each split, format [[vis1,datacolumn1,spw1,timerange1,chanwidth1,timewidth1],[vis2,datacolumn2,spw2,timerange2,chanwidth2,timewidth2],....]
	Returns
	-------
	list
		List of splited ms
	'''
	print ('Waiting for free cpus....\n')
	try:
		while True:
			available_cpu=int(psutil.cpu_count()*(1-(psutil.cpu_percent()/100.0))*float(cpu_frac)) 
			if available_cpu>0:
				print ('Number of cpus available for spliting : '+str(available_cpu)+'\n')
				break
		start_time=time.time()
		savedir=paramList[0][6]
		total_count=0
		if os.path.exists(savedir+'/temp'):
			os.system('rm -rf '+savedir+'/temp')
		os.makedirs(savedir+'/temp')
		for i in range(len(paramList)):
			t_param=paramList[i]
			t_param[-3]=savedir+'/temp'
			paramList[i]=t_param
		paramList_copy=copy.deepcopy(paramList)
		output_ms_list=[]
		verbose=paramList_copy[0][-2]
		chunksize=int(len(paramList_copy)/available_cpu)
		if chunksize<1:
			chunksize=1
		pool = mp.Pool(available_cpu) 
		ms_list = pool.map(func=split_single_ms, iterable=paramList_copy, chunksize=chunksize) 
		pool.close() 
		pool.join()
		for ms in ms_list:
			if os.path.exists(ms):
				output_ms_list.append(ms)	
			else:
				if verbose:
					print ('Msname '+ms+' does not exist.\n')
		for i in output_ms_list:	
			if os.path.exists(savedir+'/'+os.path.basename(i)):
				os.system('rm -rf '+savedir+'/'+os.path.basename(i)+'*')
			os.system('mv '+i+' '+savedir+'/')
		os.system('rm -rf '+savedir+'/temp')
		os.system('rm -rf casa*log')
		final_ms_list=glob.glob(savedir+'/*.ms')
		return 0
	except Exception as e:
		print ('Error occured during spliting : '+str(e)+'\n')
		os.system('rm -rf '+savedir+'/temp')
		os.system('rm -rf casa*log')
		return 1

def main():
	start_script_time=time.time()
	cwd=os.getcwd()
	usage= ' Split measurement set in parallel\n'
	parser = OptionParser(usage=usage)
	parser.add_option('--msname',dest="msname",default=None,help="Name of measurement set to split from",metavar="Measurement Set")
	parser.add_option('--savedir',dest='savedir',default=None,help='Name of directory to save splited measurement sets',metavar="Directory path")
	parser.add_option('--freq_interval',dest='skip_freq',default='MS frequency resolution',help='Spliting frequency interval in kHz',metavar="Float")
	parser.add_option('--time_interval',dest='skip_time',default='MS time resolution',help='Spliting temporal interval in second',metavar="Float")
	parser.add_option('--freq_list',dest='freq_list',default=None,help='Spliting frequency list in MHz (If given frequency interval will not be used)',metavar="Comma separated string")
	parser.add_option('--time_list',dest='time_list',default=None,help='Spliting timestamp list in \'yyyy/mm/dd/hh:mm:ss\' format (If given time interval will not be used)',\
						metavar="Comma separated string")
	parser.add_option('--timerange',dest='timerange',default=None,help='Spliting timerange (Data set in this timerange will be splited at the given time interval)',\
						metavar="Comma separated string")
	parser.add_option('--force',dest='force',default=False,help='Force to determine the frequency time chunks fresh even if the previous metadata exists',\
						metavar="Comma separated string")
	parser.add_option('--freq_width',dest='freq_width',default='MS frequency resolution',help='Bandwidth of each splited ms in kHz start from frequency stamp',metavar="Float")
	parser.add_option('--time_width',dest='time_width',default='MS time resolution',help='Temporal width of each splited ms in second start from time stamp',metavar="Float")
	parser.add_option('--freq_avg',dest='freq_avg',default='MS frequency resolution',help='Frequency averaging in kHz',metavar="Float")
	parser.add_option('--time_avg',dest='time_avg',default='MS time resolution',help='Temporal averaging in second',metavar="Float")
	parser.add_option('--datacolumn',dest='datacolumn',default='data',help='Datacolumn to split',metavar="String")
	parser.add_option('--cpu_frac',dest='cpu_frac',default=0.5,help='Fraction of cpu to use',metavar="Float")
	parser.add_option('--total_block',dest='total_block',default=1,help='Number of total spliting blocks',metavar="Integer")
	parser.add_option('--cur_block',dest='cur_block',default=1,help='Current spliting block',metavar="Integer")
	parser.add_option('--verbose',dest='verbose',default=False,help='Verbose output',metavar="Boolean")
	parser.add_option('--splitms',dest='splitms',default=True,help='Split ms or not',metavar="Boolean")
	parser.add_option('--max_flag_frac',dest='max_flag_frac',default=0.95,help='Maximum allowed flag fraction',metavar="Float")
	(options, args) = parser.parse_args()

	splitms=eval(str(options.splitms))
	count=0
	while True:
		available_cpu=int(psutil.cpu_count()*(1-(psutil.cpu_percent()/100.0))*float(options.cpu_frac))
		if available_cpu>0:
			break
		if count==0:
			print ('Waiting for free hardware resources....\n')
		count+=1

	if str(options.msname)[-1]=='/':
		msname=str(options.msname)[:-1]
	else:
		msname=str(options.msname)
	if os.path.isdir(str(options.msname))==False or options.msname==None:
		print ('Measurement set is not present.\n')
		os.system('touch '+cwd+'/.Finished_spliting_'+os.path.basename(str(msname))+'_noms')
		os._exit(1)
	try:
		if options.savedir==None:
			savedir=os.path.dirname(os.path.abspath(str(options.msname)))+'/splited_ms/'+os.path.basename(options.msname).split('.ms')[0]
		elif os.path.isdir(str(options.savedir)+'/splited_ms/'+os.path.basename(options.msname).split('.ms')[0])==False:
			os.makedirs(str(options.savedir)+'/splited_ms/'+os.path.basename(options.msname).split('.ms')[0])
			savedir=str(options.savedir)+'/splited_ms/'+os.path.basename(options.msname).split('.ms')[0]
		else:
			savedir=str(options.savedir)+'/splited_ms/'+os.path.basename(options.msname).split('.ms')[0]
		if os.path.isdir(savedir)==False:
			os.makedirs(savedir)
		while True:
			available_cpu=int(psutil.cpu_count()*(1-(psutil.cpu_percent()/100.0))*float(options.cpu_frac))
			if available_cpu>0:
				break
		verbose=eval(str(options.verbose))
		print ('Total '+str(available_cpu)+' cpus are available for use.\n')
		param_save_file=str(savedir)+'/'+str(os.path.basename(options.msname)).split('.ms')[0]+'_split_params'
		if eval(str(options.force))==True or os.path.exists(param_save_file+'.npy')==False:	
			AM=AccessMS(options.msname)
			if options.timerange==None and options.time_list==None:
				total_time=AM.calc_total_time()
				final_timerange=''
				if total_time>10:
					mjdtimestamps=AM.get_timestamps_in_mjdsecs()[0]
					mid_time=mjdtimestamps[int(len(mjdtimestamps)/2)]
					s_time=mid_time-5
					e_time=mid_time+5
					temp_timerange=mjdsec_to_timestamp(s_time,includedate=True,format=0) +'~'+mjdsec_to_timestamp(e_time,includedate=True,format=0) 
					timebin='10s'
				else:
					temp_timerange=''
					timebin=str(total_time)+'s'
			elif options.time_list!=None:
				user_time_list=str(options.time_list).split(',')
				temp_time_list=[]
				for t in user_time_list:
					try:
						s_time=timestamp_to_mjdsec(t.split('~')[0])
					except:
						s_time=timestamp_to_mjdsec(start_date+'/'+t.split('~')[0])
					try:
						e_time=timestamp_to_mjdsec(t.split('~')[-1])
					except:
						e_time=timestamp_to_mjdsec(start_date+'/'+t.split('~')[-1])
				temp_time_list.append(s_time)
				temp_time_list.append(e_time)
				temp_timerange=mjdsec_to_timestamp(min(temp_time_list),includedate=True,format=0)+'~'+mjdsec_to_timestamp(max(temp_time_list),includedate=True,format=0)
				fnal_timeange=copy.deepcopy(temp_timerange)
				timebin=str(max(temp_time_list)-min(temp_time_list))+'s'
			else:
				temp_timerange=str(options.timerange)
				final_timerange=str(options.timerange)
				start_time=temp_timerange.split('~')[0]
				end_time=temp_timerange.split('~')[-1]
				start_mjd=timestamp_to_mjdsec(start_time,format=0)
				end_mjd=timestamp_to_mjdsec(end_time,format=0)
				if end_time>start_time:
					timebin=str(end_mjd-start_mjd)+'s'
				else:
					timebin=''
			AM=AccessMS(str(options.msname))
			freqres=AM.calc_freqres()
			timeres=round(AM.calc_timeres(),2)
			AM=AccessMS(str(options.msname))
			unflag_freqs=AM.get_freqs()/10**6
			unflag_times_mjd=AM.get_timestamps_in_mjdsecs()[0]

			param_list=[]
			if str(options.freq_avg)=='MS frequency resolution' and str(options.time_avg)=='MS time resolution':
				chanwidth=1
				timebin=str(timeres)+'s'
			elif str(options.freq_avg)=='MS frequency resolution' and str(options.time_avg)!='MS time resolution':
				chanwidth=1
				timebin=str(options.time_avg)+'s'
			elif str(options.freq_avg)!='MS frequency resolution' and str(options.time_avg)=='MS time resolution':
				chanwidth=int(float(options.freq_avg)/freqres)
				timebin=str(timeres)+'s'
			elif str(options.freq_avg)!='MS frequency resolution' and str(options.time_avg)!='MS time resolution':
				chanwidth=int(float(options.freq_avg)/freqres)
				timebin=str(options.time_avg)+'s'
			if str(options.freq_width)=='MS frequency resolution' and str(options.time_width)=='MS time resolution':
				if str(options.skip_freq)!='MS frequency resolution':
					freq_width=float(options.skip_freq)
					skip_freq=float(options.skip_freq)/10**3
				else:
					freq_width=float(freqres/10**3)
					skip_freq=freq_width*10**3
				if str(options.skip_time)!='MS time resolution':
					time_width=float(options.skip_time)
					skip_time=float(options.skip_time)
				else:
					time_width=float(timeres)
					skip_time=float(timeres)
			elif str(options.freq_width)=='MS frequency resolution' and str(options.time_width)!='MS time resolution':
				if str(options.skip_freq)!='MS frequency resolution':
					freq_width=float(options.skip_freq)/10**3
					skip_freq=float(options.skip_freq)*10**3
				else:
					freq_width=float(freqres/10**3)
					skip_freq=freq_width*10**3
				time_width=float(options.time_width)
			elif str(options.freq_width)!='MS frequency resolution' and str(options.time_width)=='MS time resolution':
				freq_width=float(options.freq_width)/10**3
				skip_freq=freq_width*10**3
				if str(options.skip_time)!='MS time resolution':
					time_width=skip_time=float(options.skip_time)
				else:
					time_width=float(timeres)
					skip_time=time_width
			elif str(options.freq_width)!='MS frequency resolution' and str(options.time_width)!='MS time resolution':
				freq_width=float(options.freq_width)/10**3
				time_width=float(options.time_width)
				skip_freq=freq_width*10**3
				skip_time=time_width
			start_date,end_date=AM.get_obs_date()

			if options.freq_list!=None and options.time_list!=None:
				user_freq_list=str(options.freq_list).split(',')
				user_time_list=str(options.time_list).split(',')
				time_list=[]
				freq_list=[]
				for t in user_time_list:
					try:
						s_time=timestamp_to_mjdsec(t.split('~')[0])
					except:
						s_time=timestamp_to_mjdsec(start_date+'/'+t.split('~')[0])
					try:
						e_time=timestamp_to_mjdsec(t.split('~')[-1])
					except:
						e_time=timestamp_to_mjdsec(start_date+'/'+t.split('~')[-1])
					if s_time<min(unflag_times_mjd):
						s_time=min(unflag_times_mjd)
					if e_time>max(unflag_times_mjd):
						e_time=max(unflag_times_mjd)
					tm_str=mjdsec_to_timestamp(s_time,includedate=True,format=0)+'~'+mjdsec_to_timestamp(e_time,includedate=True,format=0)
					if tm_str not in time_list:
						time_list.append(tm_str)
				for f in user_freq_list:
					s_freq=float(f.split('~')[0])
					e_freq=float(f.split('~')[-1])
					s_e_freq=np.arange(s_freq,e_freq,float(skip_freq)).astype('float')
					for freq in s_e_freq:
						if freq in unflag_freqs:
							freq_list.append(freq)
			elif options.freq_list!=None and options.time_list==None:
				skip_time=int(float(skip_time)/timeres)
				user_freq_list=str(options.freq_list).split(',')
				freq_list=[]
				for f in user_freq_list:
					s_freq=float(f.split('~')[0])
					e_freq=float(f.split('~')[-1])
					s_e_freq=np.arange(s_freq,e_freq,float(skip_freq)).astype('float')
					for freq in s_e_freq:
						if freq in unflag_freqs:
							freq_list.append(freq)
				time_list=[]
				for i in range(0,len(unflag_times_mjd),skip_time):
					time_list.append(unflag_times_mjd[i])
				time_list=[mjdsec_to_timestamp(i,includedate=True,format=0) for i in time_list]
			elif options.freq_list==None and options.time_list!=None:
				skip_chan=int(float(skip_freq)/freqres)
				freq_list=[]
				for i in range(0,len(unflag_freqs),skip_chan):
					freq_list.append(unflag_freqs[i])
				user_time_list=str(options.time_list).split(',')
				time_list=[]
				for t in user_time_list:
					try:
						s_time=timestamp_to_mjdsec(t.split('~')[0])
					except:
						s_time=timestamp_to_mjdsec(start_date+'/'+t.split('~')[0])
					try:
						e_time=timestamp_to_mjdsec(t.split('~')[-1])
					except:
						e_time=timestamp_to_mjdsec(start_date+'/'+t.split('~')[-1])
					if s_time<min(unflag_times_mjd):
						s_time=min(unflag_times_mjd)
					if e_time>max(unflag_times_mjd):
						e_time=max(unflag_times_mjd)
					tm_str=mjdsec_to_timestamp(s_time,includedate=True,format=0)+'~'+mjdsec_to_timestamp(e_time,includedate=True,format=0)
					if tm_str not in time_list:
						time_list.append(tm_str)
			elif options.freq_list==None and options.time_list==None:
				skip_chan=int(float(skip_freq)/freqres)
				skip_time=int(float(skip_time)/timeres)
				freq_list=[]
				for i in range(0,len(unflag_freqs),skip_chan):
					freq_list.append(unflag_freqs[i])
				time_list=[]
				for i in range(0,len(unflag_times_mjd),skip_time):
					time_list.append(unflag_times_mjd[i])
				time_list=[mjdsec_to_timestamp(i,includedate=True,format=0) for i in time_list]
			split_spws=['0:'+str(i)+'~'+str(float(i)+float(freq_width)-(float(freqres)/10**3))+'MHz' for i in freq_list]
			if time_width>=2*timeres:
				y_string='+'+str(dttime.timedelta(seconds=time_width))
				split_timestamps=[j+y_string for j in time_list]
			else:
				split_timestamps=time_list
			split_spws=sorted(split_spws)
			split_timestamps=sorted(split_timestamps)

			for i in split_spws:
				for j in split_timestamps:
					param_list.append([str(options.msname),str(options.datacolumn),i,j,chanwidth,timebin,savedir,verbose,\
						float(options.max_flag_frac)])	
			if verbose:
				print ('Saving parameter file : '+str(param_save_file+'.npy')+'\n')
			np.save(param_save_file,param_list)
		if splitms==True:
			if verbose:
				print ('Loading spliting parameters from : '+param_save_file+'.npy\n')
			param_list=np.load(param_save_file+'.npy',allow_pickle=True)
			param_list=param_list.tolist()
			total_blocks=int(options.total_block)
			if total_blocks<1:
				total_blocks=1
			cur_block=int(options.cur_block)
			print ('Removing locks.....\n')
			AM=AccessMS(options.msname)
			AM.remove_locks()
			if total_blocks!=1:
				if int((cur_block+1)*len(param_list)/total_blocks)<len(param_list):
					param_list=param_list[int(cur_block*len(param_list)/total_blocks):int((cur_block+1)*len(param_list)/total_blocks)]
					print ('#########################################\n')
					print ('Spliting ms : '+str(options.msname)+' of block : '+str(cur_block)+' into total '+str(len(param_list))+' time frequency chuncks ..............\n')
					msg=split_ms_parallel(param_list,float(options.cpu_frac))
				elif int(cur_block*len(param_list)/total_blocks)<len(param_list):
					param_list=param_list[int(cur_block*len(param_list)/total_blocks):]
					print ('#########################################\n')
					print ('Spliting ms : '+str(options.msname)+' of block : '+str(cur_block)+' into total '+str(len(param_list))+' time frequency chuncks ..............\n')
					msg=split_ms_parallel(param_list,float(options.cpu_frac))
				else:
					print ('#########################################\n')
					print ('Nothing to split.\n')
					msg=1
			else:
				print ('#########################################\n')
				print ('Spliting ms : '+str(options.msname)+' into total '+str(len(param_list))+' time frequency chuncks ..............\n')
				msg=split_ms_parallel(param_list,float(options.cpu_frac))
			print ('Total spliting time : '+str(time.time()-start_script_time)+'s\n')	
			os.system('rm -rf casa*log')
			msdir=os.path.basename(options.msname)
			temp_ms_list=glob.glob(msdir+'/*temp.ms')
			if len(temp_ms_list)>0:
				for i in temp_ms_list:
					os.system('rm -rf '+i)
			temp_ms_list=glob.glob(savedir+'/*temp.ms')
			if len(temp_ms_list)>0:
				for i in temp_ms_list:
					os.system('rm -rf '+i)
			if os.path.exists(options.msname.split('.ms')[0]+'_temp.ms'):
				os.system('rm -rf '+options.msname.split('.ms')[0]+'_temp.ms')
			print ('#########################################\n')
			return msg
		else:
			os.system('rm -rf casa*log')
			msdir=os.path.basename(options.msname)
			temp_ms_list=glob.glob(msdir+'/*temp.ms')
			if len(temp_ms_list)>0:
				for i in temp_ms_list:
					os.system('rm -rf '+i)
			temp_ms_list=glob.glob(savedir+'/*temp.ms')
			if len(temp_ms_list)>0:
				for i in temp_ms_list:
					os.system('rm -rf '+i)
			if os.path.exists(options.msname.split('.ms')[0]+'_temp.ms'):
				os.system('rm -rf '+options.msname.split('.ms')[0]+'_temp.ms')
			print ('#########################################\n')
			return 0
	except Exception as e:
		print ('Error occured : '+str(e)+'\n')
		msdir=os.path.basename(options.msname)
		temp_ms_list=glob.glob(msdir+'/*temp.ms')
		if len(temp_ms_list)>0:
			for i in temp_ms_list:
				os.system('rm -rf '+i)
		temp_ms_list=glob.glob(savedir+'/*temp.ms')
		if len(temp_ms_list)>0:
			for i in temp_ms_list:
				os.system('rm -rf '+i)
		os.system('rm -rf casa*log')
		if os.path.exists(options.msname.split('.ms')[0]+'_temp.ms'):
			os.system('rm -rf '+options.msname.split('.ms')[0]+'_temp.ms')
		return 1

if __name__ == "__main__":
	result=main()
	os._exit(result)

