from datetime import datetime, date, time, timedelta
import pandas as pd
import numpy as np
import os
import sys
from lxml import etree
from netCDF4 import Dataset
import logging
from logging.handlers import RotatingFileHandler
import time as myTime

def main():

	J2000REF = 946679400 # 2000-01-01T00:00:00 -30min

	# CHECK ARGS

	checkArgs(3, 'Usage : python run.py (target) (Solar Wind XML config file)')

	prefix =  sys.argv[1]

	xmlFile = sys.argv[2]

	
	start_time = myTime.time()

	# GET ENVIRONMENT VARIABLES
	SW2ROOT     = os.getenv('SW2ROOT')
	SW2MISSIONS = os.getenv('REQ')
	SW2DATA     = os.getenv('SW2DATA')
	SW2NC       = os.getenv('SW2NC')

	# GET INFOS FROM XML CONFIG FILE
	start,stop,plasmaVi,srcVi,srcR,srcLon,tgtVi,tgtR,tgtLon = getXMLConfig(xmlFile)

	src = plasmaVi.split('_')[0]

	# CREATE A DIRECTORY FOR THE NEW QUERY

	MYDIR = SW2DATA+'/'+src+'/'+prefix+'_'+datetime.now().strftime('%Y%m%d')
	if not os.path.exists(MYDIR):
		os.makedirs(MYDIR,0775)

	
	# CREATION OF A LOGGER
	logFile = prefix+'_'+datetime.now().strftime('%Y%m%d')+'.log'

	logger = initLogger(MYDIR+'/'+logFile)	

	logger.info("Start : " + start)
	logger.info("Stop : " + stop)
	logger.info("Plasma VI : " + plasmaVi)
	logger.info("Source VI : " + srcVi)

	logger.info("Source distance : " + srcR)
	logger.info("Source longitude : " + srcLon)

	logger.info("Target VI : " + tgtVi)

	logger.info("Target distance : " + tgtR)
	logger.info("Target longitude : " + tgtLon)

	# COMPUTING BOUNDARIES
	marginPlasma      = getMarginDays(tgtVi) #Days
	plasmaStart       = shiftedDate(start, days=-marginPlasma, seconds=-1)
	plasmaDDStart     = time2ddtime(plasmaStart)

	# if run run.py 'alone' - for previous years
	#plasmaStop        = shiftedDate(stop, days=marginPlasma)

	plasmaStop        = shiftedDate(stop)
	plasmaDDStop      = time2ddtime(plasmaStop)
	ddTimeDeltaPlasma = DDTimeDelta(plasmaStart, plasmaStop)
	
	marginOrbit       = 20 # Hours
	orbitsStart       = shiftedDate(plasmaStart, hours=-marginOrbit)
	orbitsStop        = shiftedDate(plasmaStop, hours=marginOrbit)
	orbitsDDStart     = time2ddtime(orbitsStart)
	ddTimeDeltaOrbit  = DDTimeDelta(orbitsStart, orbitsStop)

	# GET DATA FROM DD BASE DEPENDING ON SW INPUT 
	if plasmaVi == 'omni_hour_all':
		plasmaCmd    = ['get_OMNI_1H',MYDIR+'/plasma.csv',plasmaDDStart,ddTimeDeltaPlasma]
	elif plasmaVi == 'ace_swepam_real':
		plasmaCmd    = ['get_ACE_RT',MYDIR+'/plasma.csv',MYDIR+'/mag.csv',plasmaDDStart,ddTimeDeltaPlasma]
		
		
	srcCmd   = ['get_R_LON_HCI',MYDIR+'/source.csv',srcVi,srcR,srcLon,orbitsDDStart,ddTimeDeltaOrbit]
	tgtCmd = ['get_R_LON_HCI',MYDIR+'/target.csv',tgtVi,tgtR,tgtLon,orbitsDDStart,ddTimeDeltaOrbit]
	os.system(' '.join(plasmaCmd))
	os.system(' '.join(srcCmd))
	os.system(' '.join(tgtCmd))

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

	# PLASMA DATAFRAME
	if plasmaVi == 'omni_hour_all':
		plasma = pd.read_csv(MYDIR+'/plasma.csv', dtype='S16,f4,f4,f4,f4,f4,f4,f4,f4')		
	elif plasmaVi == 'ace_swepam_real':
		if os.path.getsize(MYDIR+'/mag.csv') == 0:
			logger.error('Failed to load mag data')
			sys.exit(2)
		sw = pd.read_csv(MYDIR+'/plasma.csv', dtype='S16,f4,f4,f4')				 
		mag = pd.read_csv(MYDIR+'/mag.csv', dtype='S16,f4,f4,f4')				 
		plasma = sw.merge(mag, on='Time')
		
	plasma['Time'] = ddTime2Datetime(plasma['Time'])
	plasma = plasma.set_index('Time')
	 	 
				

	logger.info('Number of NaNs for plasma data before cleaning : %d' % plasma.isnull().sum().sum())

	plasma = plasma.interpolate().fillna(method='bfill')

	logger.info('Number of NaNs for plasma data after cleaning : %d' % plasma.isnull().sum().sum())

	# SOURCE DATAFRAME
	source = pd.read_csv(MYDIR+'/source.csv', dtype='S16,f4,f4,f4')
	source['Time'] = ddTime2Datetime(source['Time'])
	source.columns = ['Time','R_HCI_source', 'LON_HCI_source']
	source = source.set_index('Time')

	# TARGET DATAFRAME
	target = pd.read_csv(MYDIR+'/target.csv', dtype='S16,f4,f4,f4')
	target['Time'] = ddTime2Datetime(target['Time'])
	target.columns = ['Time','R_HCI_target', 'LON_HCI_target']
	target = target.set_index('Time')

	# CARTESIAN TO SPHERICAL COORDS
	r_source = source['R_HCI_source']
	lon_source_deg = np.degrees(source['LON_HCI_source'])
	lon_source_deg[lon_source_deg < 0] = lon_source_deg+360.0 #-180/180 -> 0/360

	sourceData = np.array([source.index,r_source,lon_source_deg]).T
	sourceColumns = ['Time','R_source','Lon_source']
	source = pd.DataFrame(data=sourceData,columns=sourceColumns)
	source = source.set_index('Time')

	r_target = target['R_HCI_target']
	lon_target_deg = np.degrees(target['LON_HCI_target'])
	lon_target_deg[lon_target_deg < 0] = lon_target_deg+360.0 #-180/180 -> 0/360

	targetData = np.array([target.index,r_target,lon_target_deg]).T
	targetColumns = ['Time','R_target','Lon_target']
	target = pd.DataFrame(data=targetData,columns=targetColumns)
	target = target.set_index('Time')


	# SPHERICAL TO CARTESIAN COORDS
	if plasmaVi == 'omni_hour_all':
		vlon = np.radians(plasma['Vlon'])
		vlat = np.radians(plasma['Vlat'])
	elif plasmaVi == 'ace_swepam_real': 
		vlon = np.radians(-2.0)
		vlat = np.radians(1.0)
		
	vx = plasma['V']*np.cos(vlat)*np.cos(vlon)
	vy = -plasma['V']*np.cos(vlat)*np.sin(vlon)
	vz = plasma['V']*np.sin(vlat)

	# ROTATION
	bx = -np.array(plasma['Bx'])
	by = -np.array(plasma['By'])
	bz = np.array(plasma['Bz'])

	# FINAL PLASMA DATAFRAME
	plasmaData = np.array([plasma.index,plasma.Density,plasma.Temperature,vx,vy,vz,bx,by,bz])
	plasmaColumns = ['Time','Density','Temperature','Vx','Vy','Vz','Bx','By','Bz']
	plasma = pd.DataFrame(data=plasmaData.T,columns=plasmaColumns)
	plasma = plasma.set_index('Time')

	# FINAL DATAFRAME
	combined = plasma.join(source,how='outer').join(target,how='outer')

	# trunk = combined.truncate(before=start, after=stop)
	trunk = combined.truncate(after=stop)
	df = trunk.groupby(trunk.index).first()

	df = df.resample('1H')
	df = df.interpolate().fillna(method='backfill')

	
	if df['R_target'].mean() < df['R_source'].mean():		

		df = df.reindex(index=df.index[::-1])
		idprop = -1
		xmax1  = 1.3
		xmin1  = 0.3
	else:
		idprop = 1

		xmax1  = 5.8

		xmin1  = 0.8

		
	# QUALITY FLAG
	dataFlagData = np.array( [df.index,np.zeros(len(df.index)).astype(int)])
	# dataFlagData = np.array( [df.index,np.ones(len(df.index)).astype(int)])
	dataFlagColumns = ['Time','QualityFlag']
	dataFlag = pd.DataFrame(data=dataFlagData.T,columns=dataFlagColumns)
	dataFlag = dataFlag.set_index('Time')
	df = df.join(dataFlag,how='outer')

	# WRITE TO A DATA FILE
	df.to_csv(MYDIR+'/myInputs.txt',date_format='%Y-%m-%dT%H:%M:%S',float_format="%.2f",header=False, sep=" ")
	

	# WRITE THE INPUTS FILE
	myT = [ time2J2000(str(t),J2000REF) for t in df.index ]

	filename = MYDIR+"/inputs.txt"
	fmt = 'e'
	sftFmt = ''
	with open(filename, "w") as myFile:
	    for ti,di,tempi,vxi,vyi,vzi,bxi,byi,bzi,rsrci,lonsrci,rtgti,lontgti,qfi in zip(myT, df['Density'], df['Temperature'], df['Vx'], df['Vy'], df['Vz'], df['Bx'], df['By'], df['Bz'], df['R_source'], df['Lon_source'],df['R_target'],df['Lon_target'],df['QualityFlag']):
	        line = [sftFmt,ti,format(di, fmt),format(tempi, fmt),format(vxi, fmt),format(vyi, fmt),format(vzi, fmt),format(bxi, fmt),format(byi, fmt),format(bzi, fmt),format(rsrci, fmt),format(lonsrci, fmt),format(rtgti, fmt),format(lontgti, fmt),format(qfi, fmt)]
	        myFile.write('   '.join(line)+'\n')

	myFile.close()

	# SW MHD CODE
	writeNamelist(MYDIR,idprop,xmax1,xmin1)
	logger.info('Computing solar wind propagation ...')
	os.system('sw.exe')
	
	# PROCESSING OUTPUTS OF SW CODE
	years,months,days,t,n,temp,vx,vy,by,pdyn,dphi,qualFlag = np.loadtxt(MYDIR+'/outputs.txt', comments="#", unpack="True", dtype=str,usecols=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,11))

	myDates = []
	for yi,mi,di,ti, in zip(years,months,days,t):
		if ti=='24:00:00':
			ti='23:00:00'
		myDates.append(datetime.strptime(yi+'-'+mi+'-'+di+' '+ti, '%Y-%m-%d %H:%M:%S'))
	
	myDates = np.array(myDates)

	outputsData = np.array([myDates,n.astype(float),temp.astype(float),vx.astype(float),vy.astype(float),by.astype(float),pdyn.astype(float),dphi.astype(float),qualFlag.astype(float)])
	outputsColumns = ['Time','Density','Temperature','Vx','Vy','By','pdyn','dphi','qualFlag']
	outputs = pd.DataFrame(data=outputsData.T,columns=outputsColumns)
	outputs = outputs.set_index('Time')

	outputs = outputs.truncate(before=start)

	outputs.to_csv(MYDIR+'/myOutputs.txt',date_format='%Y-%m-%dT%H:%M:%S',float_format="%.2f",header=False, sep=" ")

	# WRITE THE NETCDF OUTPUT
	ddTime = []
	for ti in outputs.index:
		ddTime.append(str(ti.year)+'{:03}'.format(ti.dayofyear-1)+'{:02}'.format(ti.hour)+'{:02}'.format(ti.minute)+'{:02}'.format(ti.second)+'000')

   	with open('ddTime.txt','w') as ddFile:
		os.system('chmod ug+w ddTime.txt')
   		for i in ddTime:
   			ddFile.write(i+'\n')
   
   	ddFile.close()

   	
	if not os.path.exists(SW2NC):
		os.makedirs(SW2NC,0775)
		
   	ncFile = SW2NC+'/'+os.path.basename(xmlFile).split('.')[0]+'.nc'
	 

	os.system('WriteDDTime '+ncFile)

	os.system('chmod ug+w ' + ncFile)

	nc = Dataset( ncFile, "a", format="NETCDF4" )

	nc.createDimension( "Vector", 2 )

	v = [ (vix, viy) for vix, viy in zip(np.array(outputs['Vx']), np.array(outputs['Vy'])) ]

	velocityVar    = nc.createVariable( 'V', 'f4', (u'Time', u'Vector') )
	velocityVar.units = 'km/s'
	velocityVar[:] = v

	magVar    = nc.createVariable( 'B', 'f4', (u'Time',) )
	magVar.units = 'nT'
	magVar[:] = np.array(outputs['By'])

	densityVar    = nc.createVariable( 'N', 'f4', (u'Time',) )
	densityVar.units = 'cm-3'
	densityVar[:] = np.array(outputs['Density'])

	temperatureVar    = nc.createVariable( 'T', 'f4', (u'Time',) )
	tempKelvin = [ tempi/11600.0 for tempi in np.array(outputs['Temperature']) ]
	temperatureVar.units = 'K'
	temperatureVar[:] = tempKelvin

	deltaPhiVar    = nc.createVariable( 'Delta_angle', 'f4', (u'Time',) )
	deltaPhiVar.units = 'degrees'
	deltaPhiVar[:] = np.array(outputs['dphi'])

	pdynVar    = nc.createVariable( 'P_dyn', 'f4', (u'Time',) )
	pdynVar.units = 'nPa'
	pdynVar[:] = np.array(outputs['pdyn'])

	qualFlagVar    = nc.createVariable( 'QualityFlag', 'i4', (u'Time',) )
	qualFlagVar[:] = np.array(outputs['qualFlag'])
	                     
	globAttr = {'Source': 'Solar Wind Model', 'Created': str(datetime.now())}
	nc.setncatts( globAttr )

	nc.close()

	# END OF PROCESS
	os.system('mv namelist '+MYDIR)
	os.system('mv fort.* '+MYDIR)
	os.system('mv ddTime.txt '+MYDIR)
	stop_time = myTime.time() 
	interval = stop_time - start_time
	logger.info('Program terminated in ' + '{:.2f}'.format(interval) + ' seconds.')

def checkArgs(nbArgs, message):
	if len(sys.argv) != nbArgs:
		print(message)
		sys.exit(2)

def initLogger(logfile):
	logger = logging.getLogger()
	logger.setLevel(logging.INFO)
 
	fileFormatter = logging.Formatter('%(asctime)s :: [%(levelname)s] :: %(message)s')
	fileHandler   = RotatingFileHandler(logfile, 'a', 1000000, 1)
	fileHandler.setLevel(logging.ERROR)
	fileHandler.setFormatter(fileFormatter)
	logger.addHandler(fileHandler)
 
	consoleFormatter = logging.Formatter('[%(levelname)s] %(message)s')
	consoleHandler   = logging.StreamHandler()
	consoleHandler.setLevel(logging.INFO)
	consoleHandler.setFormatter(consoleFormatter)
	logger.addHandler(consoleHandler)

	return logger


def getXMLConfig(XMLFilename):

	myXml    = etree.parse(XMLFilename)
	start    = myXml.xpath("/MISSION/START")[0].text
	stop     = myXml.xpath("/MISSION/STOP")[0].text
	plasmaVI = myXml.xpath("/MISSION/PLASMA_VI")[0].text
	srcVI    = myXml.xpath("/MISSION/SOURCE_VI")[0].text
	srcR     = myXml.xpath("/MISSION/SOURCE_R_PARAM")[0].text
	srcLon   = myXml.xpath("/MISSION/SOURCE_LON_PARAM")[0].text
	tgtVI    = myXml.xpath("/MISSION/TARGET_VI")[0].text
	tgtR     = myXml.xpath("/MISSION/TARGET_R_PARAM")[0].text
	tgtLon   = myXml.xpath("/MISSION/TARGET_LON_PARAM")[0].text

	return start,stop,plasmaVI,srcVI,srcR,srcLon,tgtVI,tgtR,tgtLon


def getMarginDays(tgtName):
	if tgtName == 'mercury':

		distAU == 0.61

	elif tgtName == 'venus':

		distAU = 0.28

	elif tgtName == 'mars':
		distAU = 0.52
	elif tgtName == 'jupiter':
		distAU = 4.21
	elif tgtName == 'saturn':
		distAU = 8.54
	elif tgtName == 'p67':
		distAU = 4.0
	elif tgtName == 'juno':
		distAU = 6.0
	else:
		distAU = 9.0

	return np.ceil( (distAU*150000000)/(200*24*3600) )

def shiftedDate(myDate, days=0, hours=0, minutes=0, seconds=0):
    fmt    = "%Y-%m-%dT%H:%M:%S"
    d      = datetime.strptime(myDate, fmt)

    myShiftedDate = d + timedelta(days=days, hours=hours, minutes=minutes, seconds=seconds)

    return myShiftedDate.strftime("%Y-%m-%dT%H:%M:%S")

def time2ddtime(myDate):
    fmt    = "%Y-%m-%dT%H:%M:%S"
    d      = datetime.strptime(myDate, fmt)
    dddoy  = '{:03}'.format(int(d.strftime("%j"))-1)
    hour   = '{:02}'.format(d.hour)
    minute = '{:02}'.format(d.minute)
    second = '{:02}'.format(d.second)
    
    ddtime = str(d.year) + dddoy + hour + minute + second +'000'

    return ddtime

def timeDelta(date1, date2):
    fmt    = "%Y-%m-%dT%H:%M:%S"
    d1     = datetime.strptime(date1, fmt) 
    d2     = datetime.strptime(date2, fmt) 
    td = d2 - d1

    intDays    = td.days
    intHours   = td.seconds//3600
    intMinutes = (td.seconds//60)%60
    intSeconds = td.seconds-intHours*3600-intMinutes*60

    return intDays, intHours, intMinutes, intSeconds

def DDTimeDelta(date1, date2):
    interD, interH, interM, interS = timeDelta(date1, date2)
    ddTimeDelta = '0000' + '{:03}'.format(interD) + '{:02}'.format(interH) + '{:02}'.format(interM) + '{:02}'.format(interS) + '000'

    return ddTimeDelta

def ddTime2Datetime(ddTime):
    myDatetime = []
    for t in ddTime:
        year = int(t[0:4])
        day = int(t[4:7]) + 1
        myDate = datetime(year, 1, 1) + timedelta(day - 1)
        
        hour = int(t[7:9])
        minute = int(t[9:11])
        seconds = int(t[11:13])
        ms = int(t[13:16])*1000
        myTime = time(hour, minute,seconds,ms)
        
        myDatetime.append( datetime.combine(myDate, myTime) )
            
    return myDatetime

# ISO to J2000

def time2J2000(myDate, j2000):
    fmt    = "%Y-%m-%d %H:%M:%S"
    d      = int(datetime.strptime(myDate, fmt).strftime("%s"))

    return str(d-j2000)

# make NAMELIST

def writeNamelist(directory,idprop,xmax1,xmin1):
	with open('namelist', "w") as nl:
		nl.write('&INPARA1\n')
		nl.write('	idp_in=1\n')
		nl.write('	idp_prop=1\n')
		nl.write('	idp_out=1\n')
		#nl.write('	angref(1)=0.,30.,60.,90.,120.,150.,180.,210.,240.,270.,300.,330.\n')

		nl.write('	angref(1)=0.,60.,120.,180.,240.,300.\n')

		nl.write('	idprop='+str(idprop)+'\n')
		nl.write('	fnin=\''+directory+'/inputs.txt\'\n')
		nl.write('	fnout=\''+directory+'/outputs.txt\'\n')
		nl.write('	fdirtmp=\''+directory+'/\'\n')
		nl.write('	instop=0\n')

		nl.write('	dtr=90.\n')
		nl.write('	touts=40.\n')

		nl.write('	xmin1='+str(xmin1)+'\n')
		nl.write('	xmax1='+str(xmax1)+'\n')
		nl.write('/\n')
		nl.write('&INPARA2\n')
		nl.write('	npout=10\n')
		nl.write('	bx0at1au=0.001d0\n')
		nl.write('	gm=1.40\n')
		nl.write('/\n')
	nl.close()

if __name__ == '__main__':
	main()