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
	AU = 150000000.0 # km
	
	# CHECK ARGS
	checkArgs(3, 'Usage : python run.py (target) (Solar Wind XML config file) [(SwStopTime ISO / -1)]')
	
	prefix =  sys.argv[1]
	xmlFile = sys.argv[2]
	if (len(sys.argv) == 4) :
		if ( sys.argv[3] != '-1' ) :
			swStopD = datetime.strptime(sys.argv[3], '%Y-%m-%dT%H:%M:%S.000Z')
			swStop =  swStopD.isoformat()
		else :
			swStop = '-1'
	else :
		swStop = '-1'
		
	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(prefix) #Days
	plasmaStart       = shiftedDate(start, days=-marginPlasma, seconds=-1)
	plasmaDDStart     = time2ddtime(plasmaStart)
	
	if swStop == '-1' :
		plasmaStop = shiftedDate(stop, days=marginPlasma)
	else :
		if (shiftedDate(stop, days=marginPlasma) >= swStop ) :
			plasmaStop = swStop
		else :
			plasmaStop = shiftedDate(stop, days=marginPlasma)
	# if run run.py 'alone' - for previous years	
	#plasmaStop        = shiftedDate(stop, days=marginPlasma)
	#plasmaStop        = shiftedDate(swStop)
	print plasmaStop
	 
	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]
	elif plasmaVi == 'sta_l2_pla':
		plasmaCmd    = ['get_STEREO','a',MYDIR + '/plasma.csv',MYDIR + '/mag.csv',plasmaDDStart,ddTimeDeltaPlasma]
	elif plasmaVi == 'stb_l2_pla':
		plasmaCmd    = ['get_STEREO','b',MYDIR + '/plasma.csv',MYDIR + '/mag.csv',plasmaDDStart,ddTimeDeltaPlasma]
	elif plasmaVi == 'dsc_fc_1m':
		plasmaCmd    = ['get_DSCOVR', MYDIR + '/plasma.csv',MYDIR + '/mag.csv',plasmaDDStart,ddTimeDeltaPlasma]
			
	print plasmaCmd
	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')
	elif plasmaVi == 'sta_l2_pla' or plasmaVi == 'stb_l2_pla' or plasmaVi == 'dsc_fc_1m':
		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,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'] / AU
		
	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'] / AU
	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')

	# TAO frame : RTN with t=-t
	if plasmaVi == 'sta_l2_pla' or plasmaVi == 'stb_l2_pla' :
		vx = plasma['Vr']
		vy = -plasma['Vt']
		vz = plasma['Vn']
		temperature = plasma['Vth']*plasma['Vth']*60.6 # v=sqrt(2kT/m)
	elif plasmaVi == 'dsc_fc_1m':
		vx = -plasma['Vx']
		vy = -plasma['Vy']
		vz = plasma['Vz']
		temperature = plasma.Temperature
	else :
		# 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)
		
		temperature = plasma.Temperature

	# 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,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)
	#if (prefix == 'venus' or prefix == 'mercury') :
		#trunk = combined.truncate(before=start)
	#else :
	trunk = combined #.truncate(after=stop)
		
	df = trunk.groupby(trunk.index).first()

	df = df.resample('1H')
	df = df.interpolate().fillna(method='backfill')
	
	df.Density[df.Density < 0.01] = 0.01
	
	if df['R_target'].mean() < df['R_source'].mean():		
		df = df.reindex(index=df.index[::-1])
		idprop = -1
		xmax1  = 1.3
		xmin1  = 0.3
		if prefix == 'venus':
			xmax1 = 1.5
			xmin1 = 0.5
	else:
		idprop = 1
		xmin1  = 0.8
		if prefix == 'jupiter' or prefix == 'juno' :
			xmax1 = 5.8
		elif prefix == 'mars' :
			xmax1 = 1.8
		else :
			xmax1 = 10.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:59:59'
		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, after=stop)

	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\0')

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

	nc = Dataset( ncFile, "w", format="NETCDF3_CLASSIC" )
	
	nc.createDimension( "Time", len(ddTime) )
	nc.createDimension( "TimeLength", 17 )
	nc.createDimension( "Vector", 2 )
	
	timeVar = nc.createVariable( 'Time', 'S1', ( u'Time', u'TimeLength') )
	timeStartVar = nc.createVariable( 'StartTime', 'S1', ( u'TimeLength') )
	timeStopVar = nc.createVariable( 'StopTime', 'S1', ( u'TimeLength') )
		
	timeVar[:] = np.asarray(map(list,ddTime))
	timeStartVar[:] = list(ddTime[0])
	timeStopVar[:] = list(ddTime[int(len(ddTime))-1])
	
	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 = 'eV'
	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)
	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
	elif tgtName == 'neptune':
		distAU = 29.0
	elif tgtName == 'uranus':
		distAU = 18.0	
	else:
		distAU = 9.0

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

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('	angref(1)=0.,180.\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=150.\n')
		nl.write('	touts=24.\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()