setwd("~/R/ExeterUQ/BuildEmulator")
source("DannyDevelopment.R")
source("JamesNewDevelopment.R")
library(ncdf4)

#### Load and format the data ####
#setwd("~/PhD/Canada/Wave1/MSLP")
#nc000 <- nc_open(file="psl_Amon_DevAM4-4_dma_1-000_r1i1p1_200301-200812.nc")
#ourDataPSL <- GetDataAndDesign(DataDir="", prefix="psl", variable="psl")
#ourDataPSL$FieldData <- ourDataPSL$FieldData/100 # to MB
#JJAmean <- ExtractData(OriginalField = ourDataPSL$FieldData, which.months=c(6,7,8), num.years=5)
#JJAmean$Design <- ourDataPSL$Design
#lat <- ncvar_get(nc000,"lat")
#lon <- ncvar_get(nc000,"lon")
#lon[lon>180] <- lon[lon>180] - 360
#JJAmean$lon <- lon
#JJAmean$lat <- lat
#save(JJAmean, file= "JJAmean.RData")

setwd("~/R/ExeterUQ/Demonstrations")
load("JJAmean.RData")
JJAdataPSL <- CentreAndBasis(JJAmean$Data, scaling=1)
lat <- JJAmean$lat
lon <- JJAmean$lon
plotBasis(Basis=JJAdataPSL$tBasis, 9, lat, lon, pretty(range(JJAdataPSL$tBasis[,1:9]),65))

#### Load observations ####
#tPSLFiles <- c("psl_Amon_ERA-Interim_reanalysis_r1i1p1_197901-201510.nc", "psl_Amon_MERRA_reanalysis_r1i1p1_197901-201511.nc", "psl_Amon_NCEP2_reanalysis_r1i1p1_197901-201512.nc")
#tObsJJAPSL <- GetObservations(DataFiles=tPSLFiles, which.obs="psl", EnsembleData = JJAdataPSL, permanent.scaling=100)
# Already centred
load("tObsJJAPSL.RData")

# Don't specify a weight/discrepancy, i.e. assuming a multiple of the identity, for efficiency
par(mfrow=c(1,2), mar = c(4,4,2,4))
vSVD <- VarMSEplot(JJAdataPSL, tObsJJAPSL[[1]], ylim = c(0,60))
abline(v = which(vSVD[,2] > 0.9)[1])

#### Rotate ####
JJArotPSL <- RotateBasis(JJAdataPSL, tObsJJAPSL[[1]], kmax = 5, weight = NULL, v = c(0.2,0.15,0.1,0.1,0.1), vtot = 0.9, MaxTime = 60)
vROT <- VarMSEplot(JJArotPSL, tObsJJAPSL[[1]], ylim = c(0,60))
abline(v = which(vROT[,2] > 0.9)[1])
qROT <- which(vROT[,2] > 0.9)[1]

#### Emulate ####
tDataPSL <- GetEmulatableDataWeighted(Design = JJAmean$Design, EnsembleData = JJArotPSL, HowManyBasisVectors = qROT)
emsPSL <- BasisEmulators(tDataPSL, 2, maxdf = 12)

#### Set discrepancy to avoid ruling out all of space ####
DiscMultiplier <- SetDiscrepancy(JJArotPSL$tBasis, q = qROT, obs = tObsJJAPSL[[1]])
Disc <- DiscMultiplier*diag(8192)
# To see that this does what we want:
vROTdisc <- vROT
vROTdisc[,1] <- vROT[,1] / DiscMultiplier
VarMSEplot(RecVarData = vROTdisc, obs = tObsJJAPSL[[1]], ylim = c(0,2))

#### Set a bound so can match using the coefficients ####
tempDesign <- 2*as.data.frame(randomLHS(20, 13)) - 1
colnames(tempDesign) <- colnames(JJAmean$Design)
EmOutput1 <- lapply(1:length(emsPSL), function(e) pred.em(emsPSL[[e]], tempDesign))
Expectation1 <- matrix(0, nrow = dim(tempDesign)[1], ncol = length(emsPSL))
Variance1 <- matrix(0, nrow = dim(tempDesign)[1], ncol = length(emsPSL))
for (j in 1:length(emsPSL)){
  Expectation1[,j] <- EmOutput1[[j]]$post.m
  Variance1[,j] <- EmOutput1[[j]]$post.cov
}
setwd("~/R/ExeterUQ")
Bound1 <- HMCoeffBound(JJArotPSL, tObsJJAPSL[[1]], Expectation1, Variance1, Error = 0*diag(8192), Disc = Disc)
T_c <- Bound1$Bound
stan_dens(Bound1$model, pars = "T_c")
plot(Bound1$SampleIf, Bound1$SampleIc, pch=18, xlim = c(8150,9200), ylim = c(100,800))
abline(v = qchisq(0.995, 8192))

#### Can now sample x and calculate only Ic ####
tempDesign1000 <- 2*as.data.frame(randomLHS(1000, 13)) - 1
colnames(tempDesign1000) <- colnames(JJAmean$Design)
EmOutput1000 <- lapply(1:length(emsPSL), function(e) pred.em(emsPSL[[e]], tempDesign1000))
Expectation1000 <- Variance1000 <- matrix(0, nrow = dim(tempDesign1000)[1], ncol = length(emsPSL))
for (j in 1:length(emsPSL)){
  Expectation1000[,j] <- EmOutput1000[[j]]$post.m
  Variance1000[,j] <- EmOutput1000[[j]]$post.cov
}
CoeffHM <- HistoryMatch(JJArotPSL, type="scoresMV", tObsJJAPSL[[1]], Expectation1000, Variance1000, Error = 0*diag(8192), Disc = Disc, ModelBound = FALSE)$impl
sum(CoeffHM < T_c) / length(CoeffHM)