library(isoWater)
library(assignR)
library(raster)
library(gstat)

#Get tap water data for USA
tw = wiDB_data(countries = "US", types = "Tap")
twd = tw$data

#Remove a few proprietary data
max(twd$d18O, na.rm = TRUE)
min(twd$d18O, na.rm = TRUE)
twd$d18O[twd$d18O == 9999] = NA

max(twd$d2H, na.rm = TRUE)
min(twd$d2H, na.rm = TRUE)
twd$d2H[twd$d2H == 9999] = NA

#How many records?
sum(!is.na(twd$d2H))
sum(!is.na(twd$d18O))

#Check w/ plot, looks good
plot(twd$d18O, twd$d2H)

#Remove any sites w/ no data, then make spatial
twd = twd[!(is.na(twd$d2H) & is.na(twd$d18O)),]
tw.sp = SpatialPointsDataFrame(data.frame(twd$Longitude, twd$Latitude),
                               twd[,c(1,3,4,15,16)])
proj4string(tw.sp) = "+proj=longlat +datum=WGS84 +no_defs"

#Find sites w/ same coords and average them
tw.dups = zerodist(tw.sp)
tw.dups.keep = unique(tw.dups[,1])
for(i in 1:length(tw.dups.keep)){
  kr = tw.dups.keep[i]
  rr = tw.dups[tw.dups[,1] == kr, 2]
  tw.sp$d2H[kr] = mean(tw.sp$d2H[c(kr, rr)], na.rm = TRUE)
  tw.sp$d18O[kr] = mean(tw.sp$d18O[c(kr, rr)], na.rm = TRUE)
  tw.sp$Site_ID[rr] = "remove.me"
}

tw.sp = tw.sp[tw.sp$Site_ID != "remove.me",]

#Check again
points(tw.sp$d18O, tw.sp$d2H, col = "red")

#Get precipitation grids
iso = getIsoscapes("USPrecipMA")
states.proj = spTransform(states, crs(iso))

#Project data
tw.sp = spTransform(tw.sp, crs(iso))
tw.sp = tw.sp[states.proj,]
plot(tw.sp)

#Extract values to calculate residuals
tw.sp$d2Hpcp = extract(iso$d2h_MA, tw.sp)
tw.sp$d18Opcp = extract(iso$d18o_MA, tw.sp)

#Looks at it
plot(tw.sp$d2Hpcp, tw.sp$d2H)
abline(0, 1)
plot(tw.sp$d18Opcp, tw.sp$d18O)
abline(0, 1)

#Residuals
tw.sp$d2Hres = tw.sp$d2H - tw.sp$d2Hpcp
tw.sp$d18Ores = tw.sp$d18O - tw.sp$d18Opcp

#Variograms
hvar = variogram(d2Hres ~ 1, tw.sp[!is.na(tw.sp$d2Hres),], cutoff = 2e6,
                 width = 2e6 / 20)
plot(hvar)
ovar = variogram(d18Ores ~ 1, tw.sp[!is.na(tw.sp$d18Ores),], cutoff = 2e6,
                 width = 2e6 / 20)
plot(ovar)

#Starter variogram models
hvgm = vgm(250, "Sph", 1.7e6, 80)
plot(hvar, hvgm)
ovgm = vgm(4, "Sph", 1.7e6, 2)
plot(ovar, ovgm)

#Fit them
hvgm.f = fit.variogram(hvar, hvgm, fit.method = 6)
plot(hvar, hvgm.f)

ovgm.f = fit.variogram(ovar, ovgm, fit.method = 6)
plot(ovar, ovgm.f)

#Prediction grid 25 km x 25 km
pgg = aggregate(iso$d2h_MA, 25)
pg = rasterToPolygons(pgg)

#Krige
hres = krige(d2Hres~1, tw.sp[!is.na(tw.sp$d2Hres),], newdata = pg, 
             hvgm.f, nmax = 750, maxdist = 1.5e6, debug.level = -1)
ores = krige(d18Ores~1, tw.sp[!is.na(tw.sp$d18Ores),], newdata = pg, 
             ovgm.f, nmax = 750, maxdist = 1.5e6, debug.level = -1)

#Make rasters
hres.m = rasterize(hres, pgg, hres$var1.pred)
hres.sd = rasterize(hres, pgg, sqrt(hres$var1.var))
ores.m = rasterize(ores, pgg, ores$var1.pred)
ores.sd = rasterize(ores, pgg, sqrt(ores$var1.var))

#Downscale
hres.m = resample(hres.m, iso)
hres.sd = resample(hres.sd, iso)
ores.m = resample(ores.m, iso)
ores.sd = resample(ores.sd, iso)

#Mask
hres.m = mask(hres.m, iso$d2h_MA)
hres.sd = mask(hres.sd, iso$d2h_MA)
hres.m = mask(hres.m, states.proj)
hres.sd = mask(hres.sd, states.proj)
ores.m = mask(ores.m, iso$d18o_MA)
ores.sd = mask(ores.sd, iso$d18o_MA)
ores.m = mask(ores.m, states.proj)
ores.sd = mask(ores.sd, states.proj)

#Plot
library(RColorBrewer)
par(mar = c(1,1,5,3))
cols = brewer.pal(10, "RdBu")
breaks = seq(-5, 5, length.out = 11)
breaks[11] = 5.05
plot(ores.m, axes = FALSE, breaks = breaks, col = rev(cols))
title(expression(delta^{18} * "O"["Tap Water - Precipitation"]))
lines(states.proj, col = "dark grey")
box()

#Sum
htap = iso$d2h_MA + hres.m
otap = iso$d18o_MA + ores.m

plot(otap, axes = FALSE, main = expression(delta^{18} * "O"["Tap Water"]))
lines(states.proj, col = "dark grey")

#All sample values again
tw.sp.all = SpatialPointsDataFrame(data.frame(twd$Longitude, twd$Latitude),
                               twd[,c(1,3,4,15,16)])
proj4string(tw.sp.all) = "+proj=longlat +datum=WGS84 +no_defs"
tw.sp.all = spTransform(tw.sp.all, crs(htap))
tw.sp.all = tw.sp.all[states.proj,]

#Isoscape residuals
hres.tap = tw.sp.all$d2H - extract(htap, tw.sp.all)
ores.tap = tw.sp.all$d18O - extract(otap, tw.sp.all)
hres.tap = hres.tap[!is.na(hres.tap)]
ores.tap = ores.tap[!is.na(ores.tap)]

#Compare to normal...heavy tailed
dev.off()
plot(density(hres.tap))
lines(density(rnorm(1e4, 0, sd(hres.tap))), col = "red")

#Additive variance
hres.sds = sqrt(hres.sd ^ 2 + var(hres.tap))
ores.sds = sqrt(ores.sd ^ 2 + var(ores.tap))

#Plot
plot(otap, axes = FALSE, main = expression(delta^{18} * "O"["Tap Water"]))
lines(states.proj, col = "dark grey")

#Check compatibility
stack(htap, otap, hres.sd, ores.sd, hres.sds, ores.sds)

#Write out
writeRaster(htap, "d2h.tif")
writeRaster(otap, "d18o.tif")
writeRaster(hres.sd, "d2h_se.tif")
writeRaster(ores.sd, "d18o_se.tif")
writeRaster(hres.sds, "d2h_sd.tif")
writeRaster(ores.sds, "d18o_sd.tif")