#Space/time/profiles workflow applied to SAT from Canadian model.

#Step 1 Source files and load data
source('BuildEmulator/BuildEmulator.R')
source("BuildEmulator/DannyDevelopment.R")
source("BuildEmulator/Rotation.R")
source("HistoryMatching/HistoryMatchingFast.R")
load("Demonstrations/ourDataSAT.RData")
load("Demonstrations/JJAmean.RData")
load("Demonstrations/tObsJJASAT.RData")
lat <- JJAmean$lat
lon <- JJAmean$lon

#Step 2 Extract the fields you want (This gets JJA from Canada files, new code for given netcdf formats would usually be required)
JJAmean <- ExtractData(OriginalField = ourDataSAT$FieldData, which.months=c(6,7,8), num.years=5)
JJAmean$Design <- ourDataSAT$Design

#Step 3: Centre the data and find the Singular vectors (EOFs)
JJAdataSAT <- CentreAndBasis(JJAmean$Data, scaling=1)
plotBasis(Basis=JJAdataSAT$tBasis, 9, lat, lon, pretty(range(JJAdataSAT$tBasis[,1:9]),65))

#Step 4: Set a discrepancy variance and check the performance of basis reconstructions
#Disc <- diag(runif(8192, 1, 2))# random variances
Disc <- diag(rep(0.5^2,8192))
DiscInv <- GetInverse(Disc) 
attributes(DiscInv)
par(mfrow=c(1,2), mar = c(4,4,2,4))
# Set qmax != NULL so don't calculate R_W for every truncated basis (e.g. if l is high)
vSVD <- VarMSEplot(JJAdataSAT, tObsJJASAT[[1]], weightinv = DiscInv, ylim = c(0,13.5), qmax = NULL)
abline(v = which(vSVD[,2] > 0.9)[1])

#Step 5: Rotate the basis for optimal calibration
JJArotSAT <- RotateBasis(JJAdataSAT, tObsJJASAT[[1]], kmax = 2, weightinv = DiscInv, v = c(0.25,0.15,0.1,0.1,0.1), vtot = 0.9, MaxTime = 20)
#load("Demonstrations/JJArotSAT.RData")
vROT <- VarMSEplot(JJArotSAT, tObsJJASAT[[1]], weightinv = DiscInv, ylim = c(0,13.5), qmax = NULL)
qROT <- which(vROT[,2] > 0.9)[1]
abline(v = qROT)
qROT

#Step 6: Emulate the basis coefficients and tune emulators!
tDataSAT <- GetEmulatableDataWeighted(Design = JJAmean$Design, EnsembleData = JJArotSAT, HowManyBasisVectors = qROT, weightinv = DiscInv)
tData = tDataSAT
StanEms <- InitialBasisEmulators(tData, HowManyEmulators=qROT)
tLOOs1 <- LOO.plot(StanEmulator = StanEms[[1]], ParamNames = StanEms[[1]]$Names)
newParams1 <- gpstan.default.params
newParams1$nugget <- 0.1
lmbit1 <- lapply(1:11, function(k)StanEms[[1]][[k]])
names(lmbit1) <- names(StanEms[[1]])[1:11]

StanEms[[1]] <- EMULATE.gpstan(meanResponse = lmbit1, prior.params = newParams1, tData=tData, sigmaPrior = FALSE, FastVersion = TRUE)
#Non stationary
std.err1 <- CalStError(LOOs=tLOOs1,true.y=tDataSAT$C1)
par(mfrow=c(4,4),mar=c(4,4,1,1))
for(i in 1:16){
  plot(tDataSAT[,i],std.err1,pch=20)
}
mixture1 <- MIXTURE.design(std.err1,tDataSAT[,1:13],L=2)
mixture1$MixtureMat
StanEms[[1]] <- EMULATE.gpstanNSt(meanResponse = lmbit1, prior.params = newParams1, tData=tData, sigmaPrior = FALSE, FastVersion = TRUE,mixtureComp = mixture1)
tLOOs1n <- LOO.plot.NSt(StanEmulator = StanEms[[1]], ParamNames = c(StanEms[[1]]$Names,names(StanEms[[1]]$Fouriers)))

tLOOs1 <- LOO.plot(StanEmulator = StanEms[[1]], ParamNames = StanEms[[1]]$Names)
tLOOs2 <- LOO.plot(StanEmulator = StanEms[[2]], ParamNames = StanEms[[2]]$Names)
tLOOs3 <- LOO.plot(StanEmulator = StanEms[[3]], ParamNames = StanEms[[3]]$Names)
tLOOs4 <- LOO.plot(StanEmulator = StanEms[[4]], ParamNames = StanEms[[4]]$Names)
tLOOs5 <- LOO.plot(StanEmulator = StanEms[[5]], ParamNames = StanEms[[5]]$Names)
tLOOs6 <- LOO.plot(StanEmulator = StanEms[[6]], ParamNames = StanEms[[6]]$Names)
tLOOs7 <- LOO.plot(StanEmulator = StanEms[[7]], ParamNames = StanEms[[7]]$Names)
tLOOs8 <- LOO.plot(StanEmulator = StanEms[[8]], ParamNames = StanEms[[8]]$Names)
tLOOs9 <- LOO.plot(StanEmulator = StanEms[[9]], ParamNames = StanEms[[9]]$Names)
tLOOs10 <- LOO.plot(StanEmulator = StanEms[[10]], ParamNames = StanEms[[10]]$Names)
tLOOs11 <- LOO.plot(StanEmulator = StanEms[[11]], ParamNames = StanEms[[11]]$Names)


#Adjust discrepancy if intending to history match to avoid ruling out all the input space
DiscMultiplier <- SetDiscrepancy(JJArotSAT$tBasis, q = qROT, obs = tObsJJASAT[[1]], weightinv = DiscInv, level = 0.95)
DiscScaled <- DiscMultiplier*Disc
DiscInvScaled <- (1/DiscMultiplier) * DiscInv # inverse just scales as well
par(mfrow=c(1,1))
VarMSEplot(JJArotSAT, tObsJJASAT[[1]], weightinv = DiscInvScaled, ylim = c(0,3.5), qmax = NULL)


#Step 7: History Matching
#### We can now history match large fields without needing to do large matrix inversions for every x
#### Instead, need a) reconstruction error (fixed for all x) and b) implausibility on basis, with correct projection

# Predict and history match
FieldHM <- PredictAndHM(JJArotSAT, tObsJJASAT[[1]], StanEms, tData, ns = 1000, Error = 0*diag(8192), Disc = DiscScaled, weightinv = DiscInvScaled)
summary(FieldHM$impl)
FieldHM$bound

# Plot 8 best runs + 1st ensemble member
inds <- which(FieldHM$impl <= quantile(FieldHM$impl, probs = 8 / length(FieldHM$impl))) # plot best 1%
# Order these so minimum implausibility comes first
inds <- inds[order(FieldHM$impl[inds])]
best_runs <- matrix(0, nrow = 8192, ncol = 9)
best_runs[,1] <- JJArotSAT$EnsembleMean + JJArotSAT$CentredField[,1]
for (k in inds){
  best_runs[,which(inds == k)+1] <- Reconstruct(FieldHM$Expectation[k,],JJArotSAT$tBasis[,1:qROT])+JJArotSAT$EnsembleMean
}
plotBasis(Basis=best_runs, 9, lat, lon, pretty(range(best_runs),65))

# Plot anomaly instead
best_runs_anomaly <- matrix(0, nrow = 8192, ncol = 9)
for (k in 1:9){
  best_runs_anomaly[,k] <- best_runs[,k] - JJArotSAT$EnsembleMean - tObsJJASAT[[1]]
}
max_anom <- max(abs(c(best_runs_anomaly)))
tBreaksSAT <- GenerateBreaks(WhiteRange = c(-1,1), ColRange = c(-5,5), Length.Colvec = 7,DataRange=c(-max_anom,max_anom))
colours <- function(n=21){
  intpalette(c("cyan4", "cyan1", "paleturquoise1", "white", "rosybrown2", "red2", "orangered4"), n)
}
tBreaksSAT <- seq(-15, 15, by = 2)
plotBasis(Basis=best_runs_anomaly, 9, lat, lon, tBreaksSAT, colours)

#tcoefstest <- t(Expectation1000[which.min(FieldHM),])
#CompareAnomaliesActual(x1=FieldHM$Design[which.min(FieldHM$impl),], which.member=1, Emulator=StanEms, ObsField=tObsJJASAT, EnsembleData=JJArotSAT, names.data=names(tDataSAT)[1:13], which.obs=1, AnomBreaks = tBreaksSAT, AnomCols = AnomalyColours(7), add.contours = TRUE, tcontours = seq(from=-100,by=2,to=100), tmain="SAT Anomaly [C]",cex.main=0.8,xlab="",contour.zero = FALSE)


# Create density plot
source("HistoryMatching/HistoryMatching.R")
source("HistoryMatching/impLayoutplot.R")
# Generate a larger design for this
FieldHM2 <- PredictAndHM(JJArotSAT, tObsJJASAT[[1]], StanEms, tData, ns = 10000, Error = 0*diag(8192), Disc = DiscScaled, weightinv = DiscInvScaled)
ImpData <- cbind(FieldHM2$Design, FieldHM2$impl)
ImpList <- CreateImpList(whichVars = 1:5, VarNames = colnames(tData)[1:5], ImpData, nEms=qROT, 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)