rm(list = ls())

##
## Velocity components at the center of the CME
## Figure 4 in manuscript
##

library(lubridate)
library(forecast)
library(grid)
library(gridExtra)
library(zoo)
library(date)
library(ggplot2)
library(rhdf5)
library(fields)

Sys.setenv(tz = "UTC")


vcme=1200

dcme=2

codeDir = getwd()

projectDir = dirname(codeDir)

resultsDir=paste0(projectDir,'/results')

masDir = paste0(resultsDir,'/vcme',vcme,'_dcme',dcme)

timeDir = paste0(masDir, "/dir_dump_times/")

tracerDir = paste0(masDir, "/dir_tracers/")

tracerSliceDir = paste0(masDir, "/dir_slice_tracers/")

plotsDir=paste0(projectDir,'/plots')

hdfDir= paste0(masDir, "/dir_2D/")


## MAS time units to sec

mas_to_sec = 1445.87 

MAS_to_km = 6.96 * 1e+05

m_to_au = 149597870700.

one_au = m_to_au / (MAS_to_km) / 1000.

MAS_to_km_sec = 481.3711
##
## run id details
##

yyyy0 = 2019

mm0 = 2

dd0 = 13

hh0 = 0

mm20 = 0
 
time0  = ISOdatetime(year = yyyy0, month = mm0, day = dd0, hour = hh0, min = mm20, sec = 0, tz = "UTC")



##
## read the log file of the HDF dump times for this run
## 

##
## read the log file of the HDF dumps for this run
##

setwd(masDir)

dumpDir='dir_dump_times'

setwd(dumpDir)

dump.files=list.files(path='.', pattern = 'dump_times')

ndump = length(dump.files)

list_hdf_dump_time = list_times = list_mas_df = list_date_hdf = list()

for (i in 1:ndump) {
	list_hdf_dump_time[[i]] = read.table(file=dump.files[i], header = FALSE)

}


for (i in 1:ndump) {
	##
	## need to substract relaxation time
	##
	hdf_dump_time = list_hdf_dump_time[[i]]
	
	hdf_dump_time = hdf_dump_time[, 1]
	
	deltat_ip = hdf_dump_time[1] #600.1360786600

	hdf_dump_time = hdf_dump_time - deltat_ip

	##
	## Convert to sec and add to time zero
	##

	date_hdf = hdf_dump_time * mas_to_sec + time0

	nt = length(hdf_dump_time)

	date_hdf = hdf_dump_time * mas_to_sec + time0
	
	# convert to seconds
	
	hdf_dump_time = hdf_dump_time * mas_to_sec

	# These are the HDF times
	
	times = time0 + hdf_dump_time

	list_times[[i]] = times

	list_date_hdf[[i]] = date_hdf
	
}

## 
## Cross Times
##


# Now read the tracer files for each of the ndumps

setwd(tracerDir)

## Cross Time in days
##

cross_time = times[1:ndump]
cat("\n\n Cross Times \n\n")
for (idir in 1:ndump) {
	
	dirname = paste0('tracers_wsa_z', idir)

	cat("Processing: ", dirname, '\n')
	
	setwd(dirname)

	times = list_date_hdf[[idir]]

	files = list.files(path = '.', pattern = 'tracers_pos')
	
	nfiles = length(files)
	
	fname = files[1]
	data = h5read(fname, "/Data")
	H5close()
	
	ntracers=dim(data)[1]
	
	tr = rep(NA, nfiles)
		
	for (i in 1:nfiles) {
		fname = files[i]
		data = h5read(fname, "/Data")
		tr[i] = data[1,1]	
		H5close()
	}
	
		
	index = which(tr[] >= one_au)[1]
		
	cross_time[idir] = times[index]
		
	setwd("../")	
}


dft = data.frame(dates = cross_time, variable=paste0('Realization', 1:ndump))
cat("Min/Max Arrival time: ", min(cross_time), max(cross_time), max(cross_time)-min(cross_time), '\n\n')

##
# Now read the tracer files for each of the twelve realizations and record the velocities only of the first, center of CME, particle.

setwd(tracerSliceDir)

nvars = 3

vars=c('vr', 'vt','vp','v')

for (iz in 1:ndump) {

	cat("Realization #: ", iz, "\n")

	setwd(paste0("wsa_z", iz))

	for (ivar in 1:nvars) {
		files = list.files(path = ".", pattern = paste0("tracers_", vars[ivar]))
		nfiles = length(files)
		if (ivar == 1 & iz == 1) {
			tracers = array(NA, c(nfiles, 4, ndump))
			dimnames(tracers)[[2]] = vars
		}
		cat("Processing: ", vars[ivar], "\n")
		for (j in 1:nfiles) {
			fname = files[j]
			data = h5read(fname, "/Data")
			tracers[j, ivar, iz] = data[1]
			H5close()
		}
	}

	setwd("../")
}


## Calculate |V|
for (i in 1:nfiles) {
	for (j in 1:ndump) {
		tracers[i,4,j] = sqrt(tracers[i,1,j]**2 + tracers[i,2,j]**2 + tracers[i,3,j]**2)
	}
}

# Convert to km/sec

tracers = tracers * MAS_to_km_sec


## Need to interpolate the tracers so they all run to the same clock..

time0 = round(times[1], 'hour')

timef = round(times[nfiles], 'hour')

times_1h = seq(from=time0, to=timef, by = '1 hour')

interp_tracers=tracers
## interpolate

for (i in 1:ndump) {
	x = list_times[[i]]
	x=x[1:nfiles]
	for (j in 1:4) {
		y = tracers[1:nfiles,j,i]
		interp_tracers[,j,i] = approx(x = as.numeric(x), y= y, xout = as.numeric(times_1h), rule = 2)$y 
	}
	# replace data with NA above cross time - tracer particle is at the outer BC 
	indx <- which(x > cross_time[i])
	interp_tracers[indx,1:4,i] <- NA
}

## Create a data frame for each one - easier to plot

list_df = list()

for (i in 1:4) {
	data = data.frame(dates = times_1h, interp_tracers[,i,])
	colnames(data)[2:(ndump+1)] <-  paste0('Realization', 1:ndump)
	list_df[[i]] = data
}

#dft = data.frame(dates = cross_time, variable=paste0('Realization', 1:ndump))

var.name = c("Vr [km/sec]", "Vt [km/sec]", "Vp [km/sec]", "Speed [km/sec]")

labels = c("(a)", "(b)", "(c)", "(d)")

pl = list()

for (ivar in 1:(nvars+1)) {
	df = list_df[[ivar]]
	dfc = reshape::melt(df, id.var = "dates")

	ylimits = range(dfc$value)
	pl[[ivar]] = ggplot(data = dfc, aes(x = dates, y = value, group = variable, color = variable)) + geom_line()
	
	pl[[ivar]] = pl[[ivar]] + scale_x_datetime(name = "", date_breaks = "1 day", limits = c(times_1h[1], (max(cross_time)+12*60*60)), date_labels = "%b-%d-%y")
	
	pl[[ivar]] = pl[[ivar]] + scale_y_continuous(name = var.name[ivar], limits = ylimits) + theme(text = element_text(size = 12, color = "gray20", face = "italic"), axis.text.x = element_text(face = "plain", 
		angle = 0, size = 12, vjust = 0.4), axis.text.y = element_text(face = "plain", size = 12), legend.position = "none", legend.title = element_blank())
	pl[[ivar]] = pl[[ivar]] + ggtitle(var.name[ivar]) + theme(plot.title = element_text(hjust = 0.5))

	pl[[ivar]] = pl[[ivar]] + geom_vline(xintercept = as.POSIXct(time0), color = "grey", linetype = "dashed")
	pl[[ivar]] = pl[[ivar]] + geom_vline(xintercept = as.POSIXct(time0)+(14.*60*60), color = "grey", linetype = "dashed") # End of CME
	pl[[ivar]] = pl[[ivar]] + geom_vline(data = dft, aes(xintercept=dates, group = variable, color = variable), linetype='dashed')
	pl[[ivar]] = pl[[ivar]] + annotate('text', x = (df$dates[1]), y = max(dfc$value, na.rm = TRUE), label = labels[ivar], size = 6.0) 

}
mp = grid.arrange(grobs = pl, nrow = 2, ncol = 2)


filename= paste0(plotsDir,"/figure4.pdf")

ggsave(file = filename, plot = mp, device = "pdf", width = 16, height = 10, units = "in")

cat("\n\nSaving Figure 4 in:" ,filename, '\n\n')

setwd(codeDir)