source('BuildEmulator/BuildEmulator.R')
source("BuildEmulator/DannyDevelopment.R")
source("BuildEmulator/Rotation.R")
source("HistoryMatching/HistoryMatchingFast.R")

# Select directory containing ncdf files
#setwd("ARMCU/REF")
#tt <- nc_open("SCM_1-001.nc")
#levels <- ncvar_get(tt, varid = "zf") # changes through time
#times <- ncvar_get(tt, varid = "time")
ell <- dim(levels)[1] * length(times)
#file_names <- list.files()
#extracted_data <- matrix(0, nrow = ell, ncol= length(file_names))
#for (i in 1:length(file_names)){
#  tmp_ncdf <- nc_open(file_names[i])
#  tmpdata <-  ncvar_get(tmp_ncdf, varid = "rneb")
#  extracted_data[,i] <- c(tmpdata)
#  nc_close(tmp_ncdf)
#}

# Centre the data and find the Singular vectors (EOFs)
#RNEBdata <- CentreAndBasis(extracted_data, scaling=1)
#RNEBdata$times <- times
#RNEBdata$levels <- levels
#setwd("ExeterUQ")
#save(RNEBdata, file = "Demonstrations/RNEBdata.RData")
load("Demonstrations/RNEBdata.RData")
times <- RNEBdata$times
levels <- RNEBdata$levels

# Plot basis
par(mfrow = c(3,3), mar=c(2,2,2,2))
for (i in 1:9){
  image.plot(times, (levels[91:1,1]/1000), t(matrix(RNEBdata$tBasis[,i], nrow = 91))[,91:1], ylim = c(0,4), zlim = c(-0.56,0.56))
}

# Plot mean
par(mfrow = c(1,1), mar=c(2,2,2,2))
image.plot(times, (levels[91:1,1]/1000), t(matrix(RNEBdata$EnsembleMean, nrow = 91))[,91:1], ylim = c(0,4), zlim = c(0,1),col = rainbow(n=50,start=0.5, end=1))


# Load and centre observations
obs_ncdf <- nc_open("Demonstrations/LES0.nc")
tObsRNEB <- ncvar_get(obs_ncdf, varid = "rneb")
obs_levels <- ncvar_get(obs_ncdf, varid = "zf")
image.plot(times, (obs_levels/1000), t(tObsRNEB), ylim = c(0,4), zlim = c(0,1), col = rainbow(n=50,start=0.5, end=1))


# Convert observations to same levels as data
tObsRNEBlevels <- matrix(0, nrow = 91, ncol = 16)
for (i in 1:91){
  lvl <- levels[i,1]
  closest <- which.min(abs(obs_levels - lvl))
  tObsRNEBlevels[i,] <- tObsRNEB[closest,]
}
image.plot(times, (levels[91:1,1]/1000), t(tObsRNEBlevels)[,91:1], ylim = c(0,4), zlim = c(0,1), col = rainbow(n=50,start=0.5, end=1))
tObsRNEBCentred <- c(tObsRNEBlevels) - RNEBdata$EnsembleMean
#save(tObsRNEBCentred, file = "Demonstrations/tObsRNEBCentred.RData")
#load("Demonstrations/tObsRNEBCentred.RData")

# Set a discrepancy variance and check the performance of basis reconstructions
Disc <- diag(rep(0.01^2,ell)) #### PLACEHOLDER OF 0.01
#### Discrepancy as multiple of identity works poorly. Instead downweight everything where have zeros
#diagvalues <- rep(0.01^2, ell)
#diagvalues[which(round(RNEBdata$EnsembleMean,10) == 0)] <- 1000
#diagvalues <- 1/RNEBdata$EnsembleMean
#Disc <- diag(diagvalues)

DiscInv <- GetInverse(Disc) 
attributes(DiscInv)
par(mfrow=c(1,2), mar = c(4,4,2,4))
vSVD <- VarMSEplot(RNEBdata, tObsRNEBCentred, weightinv = DiscInv)
abline(v = which(vSVD[,2] > 0.85)[1])

# Rotate the basis for optimal calibration
RNEBrot <- RotateBasis(RNEBdata, tObsRNEBCentred, kmax = 3, weightinv = DiscInv, v = c(0.45,0.15,0.1), vtot = 0.95, MaxTime = 20)
#save(RNEBrot, file = "Demonstrations/RNEBrot.RData")
#load("Demonstrations/RNEBrot.RData")
vROT <- VarMSEplot(RNEBrot, tObsRNEBCentred, weightinv = DiscInv, ylim = c(0,2))
qROT <- which(vROT[,2] > 0.85)[1]
abline(v = qROT)
qROT

par(mfrow = c(3,3), mar=c(2,2,2,2))
zmax <- max(abs(c(RNEBrot$tBasis[,1:9])))
for (i in 1:9){
  image.plot(times, (levels[91:1,1]/1000), t(matrix(RNEBrot$tBasis[,i], nrow = 91))[,91:1], ylim = c(0,4), zlim = c(-zmax,zmax))
}

# Emulate the basis coefficients and tune emulators!
load("Demonstrations/Wave1.RData")
Design <- wave_param_US[,-1]

tDataRNEB <- GetEmulatableDataWeighted(Design = Design, EnsembleData = RNEBrot, HowManyBasisVectors = qROT, weightinv = DiscInv)
tData = tDataRNEB
#tData$C1 <- log(tData$C1+150)
StanEms <- InitialBasisEmulators(tData, HowManyEmulators=1, TryFourier = FALSE, sigmaPrior = TRUE, nuggetPrior = FALSE)
for (i in 1:length(StanEms)){
  tLOOs1 <- LOO.plot(StanEmulator = StanEms[[i]], ParamNames = colnames(StanEms[[i]]$Design))
}

traceplot(StanEms[[1]]$StanModel)


# Predict and history match
FieldHM <- PredictAndHM(RNEBrot, tObsRNEBCentred, StanEms, tData, ns = 1000, Error = 0*diag(ell), Disc = Disc, weightinv = DiscInv)
summary(FieldHM$impl)
FieldHM$bound


# Plot 8 best runs + observations
inds <- which(FieldHM$impl <= quantile(FieldHM$impl, probs = 8 / length(FieldHM$impl))) # plot best runs
# Order these so minimum implausibility comes first
inds <- inds[order(FieldHM$impl[inds])]
best_runs <- matrix(0, nrow = ell, ncol = 9)
best_runs[,1] <- RNEBrot$EnsembleMean + tObsRNEBCentred # plot observations
for (k in inds){
  best_runs[,which(inds == k)+1] <- Reconstruct(FieldHM$Expectation[k,],RNEBrot$tBasis[,1:qROT])+RNEBrot$EnsembleMean
}
par(mfrow = c(3,3), mar=c(2,2,2,2))
for (i in 1:9){
  image.plot(times, (levels[91:1,1]/1000), t(matrix(best_runs[,i], nrow = 91))[,91:1], ylim = c(0,4), zlim = c(0,0.4), col = rainbow(n=50,start=0.5, end=1))
}


# Create density plot
source("HistoryMatching/HistoryMatching.R")
source("HistoryMatching/impLayoutplot.R")
# Generate a larger design for this
FieldHM2 <- PredictAndHM(RNEBrot, tObsRNEBCentred, StanEms, tData, ns = 10000, Error = 0*diag(ell), Disc = Disc, weightinv = DiscInv)
ImpData <- cbind(FieldHM2$Design, FieldHM2$impl)
ImpList <- CreateImpList(whichVars = 1:5, VarNames = colnames(tData)[1:5], ImpData, nEms=1, Resolution=c(15,15), whichMax=1, Cutoff=FieldHM2$bound)
imp.layoutm11(ImpList,VarNames = colnames(tData)[1:5],VariableDensity=TRUE,newPDF=FALSE,the.title=NULL,newPNG=FALSE,newJPEG=FALSE,newEPS=FALSE,Points=NULL)