from lxml import etree
import pandas as pd
import time as myTime

from tools import *

def main():

	# CHECK ARGS
	checkArgs(2, 'Usage : python runFootprints.py (xml config file)')
	
	if os.path.isfile(os.getenv('FOOTPRINTSMISSIONS')+'/'+sys.argv[1]) == False:
		print('[ERROR] '+sys.argv[1]+' does not exist')
		sys.exit(2)
	
	start_time = myTime.time()

	# CONNECT TO DD SERVER
	os.system('DDLogin biegun arnaud')

	# CREATE A DIRECTORY FOR THE NEW QUERY
	MYDIR = os.getenv('FOOTPRINTSDATA')+'/'+(sys.argv[1]).split('.')[0]+'_'+datetime.now().strftime('%Y%m%dT%H%M%S')
	os.system('mkdir '+MYDIR)
	os.system('chmod ug+w ' + MYDIR)
	
	# CREATION OF A LOGGER
	logFile = (sys.argv[1]).split('.')[0]+'_'+datetime.now().strftime('%Y%m%dT%H%M%S')+'.log'
	logger = initLogger(MYDIR+'/'+logFile)

	# LOAD MISSION FILE
	missionFile = os.getenv('FOOTPRINTSMISSIONS')+'/'+sys.argv[1]
	satVI, satOrbParam, satOrbUnit, plasVI, indexVI, startTime, stopTime = loadMission(missionFile)

	# GET DD TIME BOUNDARIES
	plasmaShift = 300
	plasDDStart, plasDDInter = getDDTimeBoundaries(startTime, stopTime, plasmaShift)
	
	satShift = 3600
	satDDStart, satDDInter = getDDTimeBoundaries(startTime, stopTime, satShift)
	
	indexShift = 300
	indexDDStart, indexDDInter = getDDTimeBoundaries(startTime, stopTime, indexShift)
	
	# PRINT MISSION INFOS
	logger.info("Satellite     : " + satVI)
	logger.info("Plasma        : " + plasVI)
	logger.info("Sym index     : " + indexVI)
	logger.info("Start         : " + startTime)
	logger.info("Stop          : " + stopTime)
	logger.info("Plas DD Start : " + plasDDStart)
	logger.info("Plas DD Inter : " + plasDDInter)
	logger.info("Sat DD Start  : " + satDDStart)
	logger.info("Sat DD Inter  : " + satDDInter)

	# GET DATA
	logger.info('Loading data ...')
	args        = [plasVI,plasDDStart,plasDDInter,satVI,satOrbParam,satDDStart,satDDInter,indexVI,indexDDStart,indexDDInter] 
	getDataArgs = ' '.join(args)
	os.system("GetData " + getDataArgs)

	if os.path.getsize('plasma.txt') == 0:
		logger.error('Failed to load plasma data')
		sys.exit(2)
	elif os.path.getsize('sym.txt') == 0:
		logger.error('Failed to load source orbit data')
		sys.exit(2)
	elif os.path.getsize('orbit.txt') == 0:
		logger.error('Failed to load target orbit data')
		sys.exit(2)

	os.system("mv plasma.txt "+MYDIR)
	os.system("mv sym.txt "+MYDIR)
	os.system("mv orbit.txt "+MYDIR)

	plasmaData         = pd.read_csv(MYDIR+'/plasma.txt', dtype='S16,f4,f4,f4,f4,f4,f4,f4,f4')
	plasmaData['Time'] = ddTime2Datetime(plasmaData['Time'])
	
	satData            = pd.read_csv(MYDIR+'/orbit.txt', dtype='S16,f4,f4,f4')
	satData['Time']    = ddTime2Datetime(satData['Time'])
	
	symData            = pd.read_csv(MYDIR+'/sym.txt', dtype='S16,f4')
	symData['Time']    = ddTime2Datetime(symData['Time'])

	# QUALITY FLAG FOR PLASMA DATA
	qualFlag     = np.isnan(plasmaData['vx']).astype(int)
	qualData     = {'Time': plasmaData['Time'],'QualityFlag': qualFlag}
	qualFlagData = pd.DataFrame(data=qualData)
	qualFlagData = qualFlagData.set_index('Time')

	# CLEAN PLASMA DATA
	cleanedPlasmaData = plasmaData.interpolate().fillna(method='bfill')

	# BUILD THE MAIN DATAFRAME
	myPlasma          = cleanedPlasmaData.set_index('Time')
	mySat             = satData.set_index('Time')
	mySym             = symData.set_index('Time')
	combined          = myPlasma.join(mySat,how='outer').join(mySym,how='outer').join(qualFlagData,how='outer')
	trunk             = combined.truncate(before=startTime, after=stopTime)
	df                = trunk.groupby(trunk.index).first()
	df                = df.interpolate().fillna(method='backfill')
	df                = df.resample('5min')
	df                = df[['x', 'y', 'z', 'bx', 'by','bz','vx','vy','vz','ramP','symIndex','QualityFlag']]
	df['symIndex']    = df['symIndex'].astype(int)
	df['QualityFlag'] = df['QualityFlag'].astype(int)

	# NORMALIZE SPACE STEP TO EARTH RADIUS
	if satOrbUnit != 're':
		pc = loadPCFile(os.getenv('FOOTPRINTSCONFIG')+'/physical_constants.xml')
		
		if satOrbUnit == 'km':
			df['x'] = df['x']/pc['RE']
			df['y'] = df['y']/pc['RE']
			df['z'] = df['z']/pc['RE']
		elif satOrbUnit == 'au':
			df['x'] = df['x']*(pc['AU']/pc['RE'])
			df['y'] = df['y']*(pc['AU']/pc['RE'])
			df['z'] = df['z']*(pc['AU']/pc['RE'])
		else:
			logger.error('Unit for orbit data is unknown : '+satOrbUnit)
			sys.exit(2)
			

	# WRITE INPUTS FILE
	prefix = satVI.split('_')[0] + '_' + startTime.split('-')[0]
	inputsFile = MYDIR+'/'+prefix+'.txt'
	writeInputFile(df,inputsFile)

	# COMPUTES FOOTPRINTS
	fpFile = MYDIR+'/'+prefix+'_footprints.txt'
	ncFile = MYDIR+'/'+prefix+'_footprints.nc'
	logger.info('Computation of footprints ...')
	os.system("Footprints " + inputsFile + " > " + fpFile)
	os.system("NC " + fpFile + " " + ncFile)

	# END
	stop_time = myTime.time() 
	interval = stop_time - start_time
	logger.info('Program terminated in ' + str(interval) + ' seconds.')


############################## FUNCTIONS DEFINITIONS ##################################
def loadPCFile(pcFile):
	myXml = etree.parse(pcFile)
	au    = ((myXml.xpath("/CONSTANTS/AU"))[0].text)
	re    = ((myXml.xpath("/CONSTANTS/RE"))[0].text)

	pc = {
		'AU': float(au),
		'RE': float(re),
	}
	
	return pc

def loadMission(missionFile):
	myXml       = etree.parse(missionFile)
	satName     = ((myXml.xpath("/MISSION/SATELLITE/NAME"))[0].text).lower()
	satVI       = ((myXml.xpath("/MISSION/SATELLITE/VI"))[0].text)
	satOrbParam = ((myXml.xpath("/MISSION/SATELLITE/ORBIT"))[0].text)
	satOrbUnit  = ((myXml.xpath("/MISSION/SATELLITE/UNIT"))[0].text).lower()
	plasVI      = ((myXml.xpath("/MISSION/PLASMA/VI"))[0].text)
	indexVI     = ((myXml.xpath("/MISSION/INDEX/VI"))[0].text)
	startTime   = (myXml.xpath("/MISSION/START"))[0].text
	stopTime    = (myXml.xpath("/MISSION/STOP"))[0].text

	return satVI, satOrbParam, satOrbUnit, plasVI, indexVI, startTime, stopTime

def writeInputFile(dataframe,inputsfilename):
    with open(inputsfilename, 'w') as f:
    
        for i,j in enumerate(dataframe.index):
            t = ' '.join([str(j.year),str(j.dayofyear),str(j.hour),str(j.minute)])
            x = '{:.2f}'.format(dataframe['x'][i])
            y = '{:.2f}'.format(dataframe['y'][i])
            z = '{:.2f}'.format(dataframe['z'][i])
            bx = '{:.2f}'.format(dataframe['bx'][i])
            by = '{:.2f}'.format(dataframe['by'][i])
            bz = '{:.2f}'.format(dataframe['bz'][i])
            vx = '{:.2f}'.format(dataframe['vx'][i])
            vy = '{:.2f}'.format(dataframe['vy'][i])
            vz = '{:.2f}'.format(dataframe['vz'][i])
            ramP = '{:.2f}'.format(dataframe['ramP'][i])
            sym = '{:.0f}'.format(dataframe['symIndex'][i])
            qF = '{:.0f}'.format(dataframe['QualityFlag'][i])
            params = [t,x,y,z,bx,by,bz,vx,vy,vz,ramP,sym,qF]
            line = ' '.join(params)+'\n'
            
            f.write(line)
    f.close()    
############################### END OF FUNCTIONS DEFINITIONS ############################



if __name__ == '__main__':
	main()