import os,glob,copy,time,numpy as np,matplotlib.pyplot as plt,multiprocessing as mp,logging,sys,psutil
from astropy.io import fits
from casatasks import *
from casatools import msmetadata,table,measures,quanta,agentflagger,image,calibrater,ms
from paircarstools.libpaircars import send_paircars_notification
from paircarstools.basic_func import *
from paircarstools.access_ms import *
from paircarstasks.spawn_final_imaging import final_imaging,init_logger_console
from optparse import OptionParser
from astropy import wcs
from datetime import datetime

global tracklog
formatter = logging.Formatter('%(asctime)s %(levelname)-8s %(message)s',datefmt='%Y-%m-%d %H:%M:%S')
tracklog = logging.getLogger('track_image_log')
tracklog.setLevel(logging.DEBUG)

def track_logger_console(logdir='',logger_name='',verbose=False):
	'''
	Initial logger
	Parameters
	----------
	working_dir : str
		Name of the working directory
	verbose : bool
		Verbose output or not
	'''
	if verbose==True:
		console=logging.StreamHandler(sys.stdout)
		console.setFormatter(formatter)
		tracklog.addHandler(console)
	if logdir!='' and logger_name!='':
		if os.path.exists(logdir+'/'+logger_name+'.log'):
			os.system('rm -rf '+logdir+'/'+logger_name+'.log')
		if os.path.isdir(logdir)==False:
			os.makedirs(logdir)
		filehandle=logging.FileHandler(logdir+'/'+logger_name+'.log')
		filehandle.setFormatter(formatter)
		tracklog.addHandler(filehandle)
	tracklog.propagate = False
	return

def spawn_failed_images(spawn_parameters):
	'''
	Function try to spawn imaging for failed imaging
	Parameters
	----------
	msname : str
		Name of measurement set
	basedir : str
		Name of the base directory
	savedir : str
		Name of the final image save directory
	OBSID : int
		MWA observation ID
	localdatabase : str
		Name of the local calibration database
	inputfile : str
		Name of the P-AIRCARS input file
	'''
	msname,basedir,savedir,OBSID,localdatabase,inputfile,cpu_list=spawn_parameters
	if os.path.isdir(msname)==False:
		tracklog.error('Measurement set is not present.\n')
		return 1
		
	if basedir==None:
		tracklog.error('No base directory path is given. Please give base directory path to continue.\n')
		return 1
	elif os.path.isdir(basedir)==False:
		os.makedirs(basedir)

	if inputfile==None or os.path.isfile(inputfile)==False:
		tracklog.error('P-AIRCARS input file is not present. Please give the input file and re run.\n')
		return 1

	os.chdir(basedir)
	if inputfile[-1]=='/':
		inputfile=inputfile[:-1]
	cwd=os.getcwd()
	if os.path.isfile(cwd+'/selfcal_inputs.py')==False and (inputfile==None or os.path.isfile(inputfile)==False): 
		tracklog.error('Input file does not exist.1\n')
		return 1
	elif inputfile!=None and os.path.isfile(cwd+'/selfcal_inputs.py')==False:
		os.system('cp -r '+inputfile+' '+cwd+'/selfcal_inputs.py')
	inputfile=cwd+'/selfcal_inputs.py'
	sys.path.append(cwd)
	import selfcal_inputs as inputs

	workdir=os.path.dirname(os.path.abspath(msname))+'/workdir_'+os.path.basename(msname).split('.ms')[0]
	AM=AccessMS(str(msname))
	flagfrac=AM.calc_flagfrac()
	if flagfrac>=0.9:
		tracklog.error('Flag fraction is more than 90%.\n')
		imagedir=savedir+'/All_final_images/'+str(OBSID)
		file_str=workdir+'/'+os.path.basename(splited_ms_rename(msname,ref_time_chan=False,change_msname=False)).split('.ms')[0]
		os.system('touch '+imagedir+'/'+os.path.basename(file_str)+'_image.junk')
		os.system('rm -rf '+msname+' '+msname+'.flagversions')
		return 1
	freqs=AM.get_freqs()/10**6
	bandwidth=AM.calc_bandwidth()/10**3
	timewidth=AM.calc_total_time()
	coarse_chan_0=freq_to_MWA_coarse(freqs[0])
	coarse_chan_1=freq_to_MWA_coarse(freqs[-1])
	is_internet=check_internet()
	metafits=glob.glob(basedir+'/data/'+str(OBSID)+'.metafits')
	if len(metafits)==0:
		metafits=download_metafits(msname,basedir+'/data')
	else:
		metafits=metafits[0]

	if inputs.keep_logger:
		logdir=str(basedir)+'/verbose_logs/'+str(OBSID)
		logger_name=str(OBSID)+'_final_imagelog'	
	else:
		logdir=logger_name=''
	init_logger_console(logdir=logdir,logger_name=logger_name,verbose=inputs.verbose)
	caltables=glob.glob(str(localdatabase)+'/'+str(OBSID)+'/'+str(OBSID)+'*_'+str(coarse_chan_0)+'*cal')
	if len(caltables)==0:
		tracklog.error('No caltables are present in local caldatabase : '+str(localdatabase)+'/'+str(OBSID)+'\n')
		return 1

	if inputs.calc_selfcalib_params==True:
		calparams=CalcParams(msname,inputs.quality_factor,inputs.safety_factor)
		residual_frac=float(calparams.calc_calibration_params()[2])
	else:
		residual_frac=0.1

	if inputs.do_polcal==True:
		do_polcal=False	
		for cal in caltables:
			if '.pcal' in cal:
				do_polcal=True
		if do_polcal==False:
			tracklog.warning('Polarization caltables are not present. Hence not performing polarization imaging.\n')
	else:
		do_polcal=False
		tracklog.info('User wants only Stokes I imaging.\n')

	if inputs.maskfile!='':
		mask=inputs.maskfile
	elif inputs.maskstr!='':
		mask=inputs.maskstr
	else:
		mask=''
	try:
		rms_list=np.load(basedir+'/Ref_time_chan_sigma.npy',allow_pickle=True)[1]
	except Exception as e:
		tracklog.warning('Reference time rms threshold could not be found. Setting it to default value.\n')
		rms_list=[0.1]
	if float(inputs.image_time)>2:
		timeavg=2
	else:
		timeavg=float(inputs.image_time)
	if float(inputs.image_freq)>160:
		freqavg=160
	else:
		freqavg=float(inputs.image_freq)
	if do_polcal==True:
		stokes='IQUV'
		if len(rms_list)!=4:
			thresh=rms_list*4
	else:
		if len(rms_list)>1:
			thresh=rms_list[0]
		stokes='pseudoI'
	thresh=[str(i) for i in thresh]
	threshold=','.join(thresh)
	imaging_mode='final'
	caltable_str=''
	for cal in caltables:
		caltable_str+=cal+','
	caltable_str=caltable_str[:-1]
	cmd_args=['--msname '+msname,' --metafits '+metafits,' --basedir '+basedir,' --workdir '+workdir,' --savedir '+savedir,' --savemodel '+str(inputs.savemodel),\
		' --saveres '+str(inputs.saveresidual),' --stokes '+stokes,' --sigma 6',' --threshold '+str(threshold),' --wsclean '+str(inputs.use_wsclean),' --verbose '+str(inputs.verbose),\
		' --want_automask '+str(inputs.want_auto_masking),' --quality_factor '+str(inputs.quality_factor),' --use_ankflag '+str(inputs.use_ankflagger),' --imaging_mode '+str(imaging_mode),\
		' --residual_frac '+str(residual_frac),'--caltables '+str(caltable_str)]#,'--major_axis '+str(maj),'--minor_axis '+str(minor),'--pa '+str(pa)]
	cmd_args.append(' --maskfile '+str(mask))
	cmd_args.append(' --cutoutbox '+inputs.cutoutbox)
	cmd_args.append(' --num_pixel_in_psf '+str(5))
	cmd_args.append(' --time_avg '+str(timeavg))
	cmd_args.append(' --freq_avg '+str(freqavg))
	cmd='paircars_final_imaging '+' '.join(cmd_args)
	screen_name=str(OBSID)+'_'+imaging_mode+'_'+os.path.basename(msname).split('.ms')[0]+'_'+str(inputs.job_id)
	finished_touch_file=basedir+'/.Finished_final_imaging_'+str(imaging_mode)+'_'+os.path.basename(msname).split('.ms')[0]
	batch_file=paircars_instance_runner(cmd,basedir,inputs.paircars_dir,screen_name,finished_touch_file,inputs.job_id)
	run_batch_scripts_cpulist(batch_file,cpu_list=cpu_list)
	return finished_touch_file

def track_imaging(obsid,basedir,savedir,expected_image_num,localdatabase,email):
	'''
	Function to track final imaging
	Parameters
	----------
	obsid : int
		Observation ID		
	basedir : str
		Name of the base directory
	savedir : str
		Name of the save directory
	expected_image_num : int
		Approximate image to be expected
	localdatabase : str
		Path of the local calibration database
	email : str
		E-mail address to send notification
	'''	
	tracklog.info('Waiting for final images.....\n')
	while True:
		time.sleep(0.5)
		if os.path.isdir(savedir+'/All_final_images/'+str(obsid)):
			os.chdir(savedir+'/All_final_images/'+str(obsid))
			break
	old_finished_num=0
	imaging_finished=False
	is_internet=check_internet()
	update_time=0
	failed_imaging=True
	while imaging_finished==False:
		#try:
		time.sleep(10)
		imagelist=glob.glob('*.fits')+glob.glob('*.junk')
		total_final_images=0
		finished_imaging_touch_files=glob.glob(basedir+'/.Imaging_done_*')
		for i in finished_imaging_touch_files:
			total_final_images+=int(i.split('_')[-1])
		if len(imagelist)>=expected_image_num or len(imagelist)>=total_final_images:
			successful_images=glob.glob('*.fits')
			failed_images=glob.glob('*.junk')
			imaging_finished=True
		else:
			successful_images=glob.glob('*.fits')
			failed_images=glob.glob('*.junk')
			imaging_finished=False
		print (total_final_images,expected_image_num,len(imagelist))
		if imaging_finished:
			try_count=0
			while True:
				failed_ms=glob.glob(savedir+'/failed_ms/'+str(obsid)+'/*.ms')
				failed_ms_copy=copy.deepcopy(failed_ms)
				for i in failed_ms_copy:
					if 'cor' in i or 'avg' in i or 'flagversions' in i:
						failed_ms.remove(i)
				del failed_ms_copy
				if len(failed_ms)==0:
					break
				else:
					inputfile=basedir+'/selfcal_inputs.py'
					cpu_list_copy=np.load(basedir+'/cpulist.npy',allow_pickle=True).tolist()
					cpu_frac=len(cpu_list_copy)/psutil.cpu_count()	
					free_cpus=get_free_cpu_list(cpu_frac,per_cpu_free_frac=0.25,cpu_list=cpu_list_copy)	
					paircars_instance=int(len(free_cpus)/3)
					free_cpus_copy=copy.deepcopy(free_cpus)
					tracklog.info('Re-trying to make final images of the measurement sets : '+','.join([os.path.basename(i) for i in failed_ms]))
					params=[]
					for msname in failed_ms:
						if len(free_cpus)==0:
							free_cpus=copy.deepcopy(free_cpus_copy)
						cpu_list=free_cpus[:3]
						for i in cpu_list:
							free_cpus.remove(i)
						params.append([msname,basedir,savedir,obsid,localdatabase,inputfile,cpu_list])
					chunksize=int(len(params)/paircars_instance)
					if chunksize<1:
						chunksize=1
					imaging_done=0
					prev_imaging_done=0
					results=[]
					while len(params)>0:
						pool = mp.Pool(chunksize)
						temp_params=params[0:chunksize]	
						for i in range(len(temp_params)):
							tracklog.info('spawn_failed_images('+str(temp_params[i])+')\n')
						temp_results = pool.map(func=spawn_failed_images, iterable=temp_params, chunksize=1) 	
						for p in temp_params:
							params.remove(p)
						for i in temp_results:
							results.append(i)		
						pool.close() 
						pool.join()
						c=0 
						if len(params)==0:
							print ('Imaging have been spawned for all measurement sets.\n')
							break
						while True:
							for r in results:
								if r==1:
									imaging_done+=1
								elif len(glob.glob(r+'*'))>0:
									imaging_done+=1
							if imaging_done-prev_imaging_done>1:
								chunksize=imaging_done
								prev_imaging_done=imaging_done
								print ('Imaging have been done for '+str(imaging_done)+' ms.\n')
								break
							else:
								if c==0:
									print ('Waiting for finishing imaging..\n')
								c+=1
								time.sleep(5)
					try_count+=1
					c=0 
					imaging_done=0
					time.sleep(10)
					while True:
						for r in results:
							if r==1:
								imaging_done+=1
							elif len(glob.glob(r+'*'))>0:
								imaging_done+=1
						if imaging_done>=len(results):
							print ('Imaging finished for all measurement sets.\n')
							break
						else:
							if c==0:
								print ('Waiting for finishing imaging..\n')
							c+=1
							time.sleep(5)
					if try_count==10:
						print ('P-AIRCARS tries 10 times to make final images. There are still some images could not be made for measurement sets : '+','.join([os.path.basename(i) for i in failed_ms]))
			
		if imaging_finished==False:
			tracklog.info('Total successful images made till now : '+str(len(successful_images))+'\n')
			tracklog.info('Total failed images made till now : '+str(len(failed_images))+'\n')
			present_time=time.time()
			if is_internet and email!=None and (update_time==0 or (present_time-update_time)>600):
				imagefiles = glob.glob('time*.fits')
				imagefiles.sort(key=lambda x: os.path.getmtime(x))
				if len(imagefiles)>0:
					imagename=imagefiles[-1]
				else:
					imagename=''
				if len(successful_images)>old_finished_num:
					old_finished_num=len(successful_images)
					update_time=present_time
					try:
						if imagename=='':
							msg_str='Dear P-AIRCARS user,\n\nFinal imaging for OBSID : '+str(obsid)+' is running.\n Number of images made : '+str(len(successful_images))+\
										'.\n\nP-AIRCARS '
						else:
							msg_str='Dear P-AIRCARS user,\n\nFinal imaging for OBSID : '+str(obsid)+' is running.\n Number of images made : '+str(len(successful_images))+\
										'.\nLast image : '+os.path.basename(imagename)+'\n\nP-AIRCARS '
						msg_subject='Notification from P-AIRCARS : Final imaging : OBSID = '+str(obsid)
						send_paircars_notification(email,msg_subject,msg_str,attachments=[])
					except Exception as e:
						tracklog.warning('Error : '+str(e)+'\n')
						pass
		else:
			tracklog.info('Total successful images : '+str(len(successful_images))+'\n')
			tracklog.info('Total failed images : '+str(len(failed_images))+'\n')
			present_time=time.time()
			if is_internet and email!=None and (update_time==0 or (present_time-update_time)>600):
				if len(successful_images)>old_finished_num:
					old_finished_num=len(successful_images)
					update_time=present_time
					try:
						imagefiles = glob.glob('time*.fits')
						imagefiles.sort(key=lambda x: os.path.getmtime(x))
						if len(imagefiles)>0:
							imagename=imagefiles[-1]
						else:
							imagename=''
						if imagename=='':
							msg_str='Dear P-AIRCARS user,\n\nFinal imaging for OBSID : '+str(obsid)+' is running.\n Number of images made : '+str(len(successful_images))+\
										'.\n\nP-AIRCARS '
						else:
							msg_str='Dear P-AIRCARS user,\n\nFinal imaging for OBSID : '+str(obsid)+' is running.\n Number of images made : '+str(len(successful_images))+\
										'.\nLast image : '+os.path.basename(imagename)+'\n\nP-AIRCARS'
						msg_subject='Notification from P-AIRCARS : Final imaging : OBSID = '+str(obsid)
						send_paircars_notification(email,msg_subject,msg_str,attachments=[])
					except Exception as e :
						tracklog.warning('Error : '+str(e)+'\n')
			tracklog.info('#########################\nFinal imaging for : '+str(obsid)+' is finished.\nTotal number of images made : '+str(len(imagelist))+'\n#########################\n')
			return
			break	
		'''except Exception as e:
			tracklog.error('Errror in sending imaging notification : '+str(e)+'\n')
			pass'''



def send_database_images(obsid,basedir,localdatabase):
	'''
	Function to send database images to database
	Parameters
	----------
	obsid : int
		Observation ID
	basedir : str
		Name of the base directory
	localdatabase : str
		Path of the local database		
	'''
	print ('Waiting for database images.....\n')
	while True:
		time.sleep(180)
		if os.path.isdir(localdatabase+'/'+str(obsid)+'/images/All_final_images/'+str(obsid)):
			os.chdir(localdatabase+'/'+str(obsid)+'/images/All_final_images/'+str(obsid))
			break
	old_finished_num=0
	imaging_finished=False
	is_internet=check_internet()
	cwd=os.getcwd()
	while imaging_finished==True:
		time.sleep(600)
		imagelist=glob.glob('*.fits')+glob.glob('*.junk')
		total_final_images=0
		finished_imaging_touch_files=glob.glob(basedir+'/.Database_imaging_done_*')
		for i in finished_imaging_touch_files:
			total_final_images+=int(i.split('_')[-1])
		if len(imagelist)>=total_final_images:
			successful_images=glob.glob('*.fits')
			failed_images=glob.glob('*.junk')
			imaging_finished=True
		else:
			successful_images=glob.glob('*.fits')
			failed_images=glob.glob('*.junk')
			imaging_finished=False
		if imaging_finished:
			print ('Total successfull database images : '+str(len(successful_images))+'\n')
			print ('Total failed images : '+str(len(failed_images))+'\n')
			if is_internet:
				zip_file_name=str(obsid)+'_database_images.zip'
				with zipfile.ZipFile(zip_file_name,mode='w') as zf:
					for image in successful_images:
						zf.write(os.path.basename(image))
				msg_str='Dear P-AIRCARS User,\n\nDatabase imaging for OBSID : '+str(obsid)+' is done.\nTotal images made : '\
							+str(len(successful_images))+'\n\nP-AIRCARS '
				msg_subject='Notification from PAIRCARS : Database imaging : OBSID = '+str(obsid)
				send_paircars_notification('paircarsnotification@gmail.com',msg_subject,msg_str,attachments=[zip_file_name])
				database_notification_sent=1
				os.system('rm -rf '+zip_file_name)
				os.system('rm -rf '+localdatabase+'/'+str(obsid)+'/images/All_final_images/'+str(obsid)+'/*.junk')
				os.system('rm -rf '+localdatabase+'/'+str(obsid)+'/images/splited_ms/*')
				os.chdir(cwd)
			break
		os.chdir(cwd)
			

def main():
	usage= 'Track and notifier for final imaging'
	parser = OptionParser(usage=usage)
	parser.add_option('--OBSID',dest="obsid",default=0,help="Observation ID",metavar="Integer")
	parser.add_option('--basedir',dest="basedir",default=None,help="Name of the base directory",metavar="Directory name")
	parser.add_option('--localdatabase',dest="localdatabase",default=None,help="Name of the local database",metavar="Directory name")
	parser.add_option('--savedir',dest="savedir",default=None,help="Name of the directory where images are saved",metavar="Directory name")
	parser.add_option('--email',dest="email",default=None,help="E-mail ID to send notification",metavar="E-mail ID")
	parser.add_option('--database_imaging',dest="database_imaging",default=True,help="Database imaging",metavar="Boolean")
	parser.add_option('--only_cal',dest="only_cal",default=False,help="Only performing calibration",metavar="Boolean")
	parser.add_option('--expected_image_num',dest="expected_image_num",default=1,help="Approximate image number to be expected",metavar="Integer")
	(options, args) = parser.parse_args()
	if options.basedir==None or os.path.isdir(options.basedir)==False:
		print ('Base directory does not exists\n')
		return 1
	if options.savedir==None or os.path.isdir(options.savedir)==False:
		print ('Save directory does not exists\n')
		return 1	
	if int(options.obsid)<=0:
		print ('Please provide correct observation ID.\n')
		return 1
	start_time=time.time()
	now=datetime.now()	
	dt_string=now.strftime("%Y-%m-%d %H:%M:%S")
	track_logger_console(logdir='',logger_name='',verbose=True)
	track_imaging(options.obsid,options.basedir,options.savedir,int(options.expected_image_num),options.localdatabase,options.email)
	return 0

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