source('lib.R')

load_packages(c('parallel','rstan','evd','data.table',
    'reshape2','SpatialExtremes','ggplot2','ismev','dplyr',
    'fields','rgdal','sp','tools','akima',
    'maps','boot','mvtnorm','digest'),quietly=TRUE)

#set_cppo('fast')

    # sampler iterations
iterations = 3500
warmup = 1500
group_size = 30
season_ = 'fall'
no_nugget = FALSE
nchains = 1
n_beta_knots = 10
# dont fit the model, only read in the complete model and postprocess or plot
postprocess_only = TRUE
plot_only = TRUE
max_treedepth = 10
# lower target acceptance rate to get a larger stepsize 
# and fewer leapfrog steps
adapt_delta = 0.6
# use these to set maximum nubers of stations for missing and complete groups
# it is useful when you want the same number of stations while changing the 
# group size. If this is not necessary set max_m/c to be very large.
#max_m = 180
#max_c = 60
max_m = 9999
max_c = 9999
validation = FALSE
percent_drop = 0.5


model_name_ = 'gev_spatial_composite_beta'
model_file = "gev_spatial_composite_beta.stan"
plot_dir = 'plots'
# the real model name
model_name = sprintf('%s_%s_%02dg_%02dknots_td%02d_ad%4.2f',model_name_,season_,group_size,n_beta_knots,max_treedepth,adapt_delta)
#model_name = sprintf('%s_%s_%02dg_%02dknots_td%02d',model_name_,season_,group_size,n_beta_knots,max_treedepth)
message('This model is: ',model_name)

############################################################
# output files
output_dir = sprintf('output/%s',model_name)
plot_dir = sprintf('%s/plots/',output_dir)
sample_dir = sprintf('%s/samples/',output_dir)
sample_file = sprintf('%s/samples',sample_dir)
diagnostic_file = sprintf('%s/diagnostics',sample_dir)
posterior_plot_dir = sprintf('%s/posterior_plots/',plot_dir)
trace_plot_dir = sprintf('%s/trace_plots/',plot_dir)
model_plot_dir = sprintf('%s/model_plots/',plot_dir)
posterior_data = sprintf('%s/stanfit_%s.rds',output_dir,model_name)
grid_data = sprintf('%s/grid_%s.rds',output_dir,model_name)
quantile_data = sprintf('%s/quantiles_%s.rds',output_dir,model_name)
dir_create(c(posterior_plot_dir,trace_plot_dir,model_plot_dir,sample_dir))
############################################################

############################################################
# seed for the random number generator so that 
# the results are more reproducible
rng_seed = 7777
init_rng_seed = 7777
nsamples = ((iterations-warmup)*nchains)
############################################################

############################################################
# Data and covariates
map_complete = fread('data/mean_daily_precip_with_elev.csv')
# observed data
#long = fread('data/complete_1960_2013.csv')
long = fread('data/west_us_precip_3day_max.csv')[year>=1960 & year<=2013 & max<10000]
# & lon>-125 & lon < -120 & lat >40 & lat<45]
long[,n:=length(max),by=c('station','season')]
# convert to cm
long[,max:=max/100]
maxn = max(long$n)

long_season = long[season==season_ & n>=30]
stns_ = unique(long_season,by='station')
map = map_complete[season==season_]
missing_elev = which(is.na(with(na.omit(map),interpp(lon,lat,elev,stns_$lon,stns_$lat)$z)))
missing_map = which(is.na(with(map,interpp(lon,lat,map,stns_$lon,stns_$lat)$z)))
missing_stns = stns_$station[unique(c(missing_elev,missing_map))]
long_season = long_season[!(station %in% missing_stns)]
long_mean = summarise(group_by(long_season,station,lon,lat),mean=mean(max))

if(validation){
    set.seed(rng_seed)
    missing_data = summarise(group_by(long_season,station), missing = !all(1960:2013 %in% year))
    n_missing = nrow(missing_data[missing == TRUE])
    n_drop = floor(percent_drop * n_missing)
    message('Dropping ', n_drop, ' stations for cross validation.')
    drop_stations = sample(missing_data[missing == TRUE]$station,n_drop)
    long_season = long_season[!(station %in% drop_stations)]
}

# data matrix 
y_ = dcast(long_season,year~station,value.var='max')
y_$year = NULL

complete = sapply(y_,function(x)!any(is.na(x)))
complete_stns = names(y_)[complete]
missing_stns = names(y_)[!complete]

# drop some stations so that there is an equal number in each group
# also randomize the station order
S_c_ = length(complete_stns)
S_m_ = length(missing_stns)
G_c = min(floor(S_c_/group_size), floor(max_c/group_size))
G_m = min(floor(S_m_/group_size), floor(max_m/group_size))
n_drop_c = S_c_ - G_c*group_size
n_drop_m = S_m_ - G_m*group_size

message('Dropping ', n_drop_c, ' complete stations.')
message('Dropping ', n_drop_m, ' incomplete stations.')

# drop some stations randomly from those with the lowest 20% mean 
set.seed(rng_seed)
drop_c = sample(long_mean[station %in% complete_stns][order(mean)[1:max(0.2*S_c_,n_drop_c)]]$station,n_drop_c)
drop_m = sample(long_mean[station %in% missing_stns][order(mean)[1:max(0.2*S_m_,n_drop_m)]]$station,n_drop_m)


complete_station_order = sample(complete_stns[!(complete_stns %in% drop_c)])
missing_station_order = sample(missing_stns[!(missing_stns %in% drop_m)])

long_season_c = long_season[station %in% complete_station_order]
long_season_m = long_season[station %in% missing_station_order]

long_combined = rbind(long_season_c,long_season_m)

# data matrix 
y_c = dcast(long_season_c,year~station,value.var='max')[,complete_station_order]
y_c$year = NULL
y_m = dcast(long_season_m,year~station,value.var='max')[,missing_station_order]
y_m$year = NULL

y = cbind(y_c,y_m)

S_c = ncol(y_c)
S_m = ncol(y_m)
T = nrow(y)
S = ncol(y)
G = floor(S/group_size)

message('Using ', S_c, ' complete stations.')
message('Using ', S_m, ' incomplete stations.')
message('Using ', S_c+S_m, ' total stations.')

stns = unique(long_combined,by=c('station'))

r_lon = c(extendrange(stns$lon)[1],range(stns$lon)[2])
r_lat = extendrange(stns$lat)
map_region = map[lon>=r_lon[1]&lon<=r_lon[2]&lat>=r_lat[1]&lat<=r_lat[2]]
d_grid = data.table(project(as.matrix(map_region[,c('lon','lat'),with=F]),
    proj="+proj=utm +ellps=WGS84 +zone=21 +units=m")/1000/1000) #Mm
setnames(d_grid,c('x','y'))
map_region = cbind(map_region,d_grid)

# covariates, can incorporate others here 
ll = unique(long_combined,by=c('lon','lat'))[,c('lon','lat'),with=FALSE]
d = project(as.matrix(ll),proj="+proj=utm +ellps=WGS84 +zone=21 +units=m")/1000/1000
dist = as.matrix(dist(d))

#spatial_covariates = as.matrix(unique(long_complete,by='station')[,'elevation',with=F])
# add spatial coordinates as covariates, could use squared cooords here
spatial_covariates = cbind(
    with(map,interpp(lon,lat,map,stns$lon,stns$lat))$z, # MAP
    with(na.omit(map),interpp(lon,lat,elev,stns$lon,stns$lat))$z) # Elevation
    #) # lat, lon

covars = scale(spatial_covariates)
class(covars) <- "trend.spatial"

P = ncol(covars)
#if(S %% G != 0) G = G + 1

y_mat = matrix(as.numeric(as.matrix(y)),T) 
############################################################

############################################################
# knot locations
cover = cover_design_cache(d_grid,n_beta_knots)
knot_ind = cover$best.id
knots = d_grid[knot_ind,]

# distance from knots to other knots 
dist_knots_to_points = rdist(d,knots)
dist_knots_to_grid = rdist(d_grid,knots)
#knot_dist = as.matrix(dist(knots))
############################################################

############################################################
# MLE estimates of latent parameter processes
# fit a gev to each station to get initial estimate of parameters
fits = suppressWarnings(lapply(y,gev_mle_fit,ignore_zeros=TRUE,maxit=1e6,method="L-BFGS-B"))
mle_loc = sapply(fits,'[',1)
mle_scale = sapply(fits,'[',2)
mle_shape = sapply(fits,'[',3)
nofit = mle_shape == -99
for(i in which(nofit)){
    mle_loc[i] <- mean(mle_loc[mle_loc!=-99],na.rm=T)
    mle_scale[i] <- mean(mle_scale[mle_scale!=-99],na.rm=T)
    mle_shape[i] <- .1
}
mle_shape[mle_shape>0.5] <- 0.49
mle_shape[mle_shape< -0.5] <- -0.49

gev_mle = as.data.table(cbind(mle_loc,mle_scale,mle_shape))
setnames(gev_mle,c('loc','scale','shape'))
gev_mle[,station:=names(mle_loc)]
gev_mle = merge(gev_mle,stns,by='station')[station %in% c(complete_station_order,missing_station_order)]


coef_loc = coefficients(lm(mle_loc~covars+0))
coef_scale = coefficients(lm(mle_scale~covars+0))
coef_shape = coefficients(lm(mle_shape~covars+0))

gev_mle[,loc_mean:=covars %*% coef_loc]
gev_mle[,scale_mean:=covars %*% coef_scale]
gev_mle[,shape_mean:=covars %*% coef_shape]

gev_mle[,loc_resid:=loc-loc_mean]
gev_mle[,scale_resid:=scale-scale_mean]
gev_mle[,shape_resid:=shape-shape_mean]

############################################################
# define design matricies 
#covars_pred = data.table(cbind(1,scale(map_region[,c('map','elev','x','y'),with=F])))
#setnames(covars_pred, c('intercept','map','elevation','x','y'))

covars_pred = data.table(scale(map_region[,c('map','elev'),with=F]))
setnames(covars_pred, c('map','elevation'))
covars_pred = as.matrix(covars_pred)
############################################################

#############################################################
#gp_fit_loc =   likfit_cache(mle_loc,  d,covars, nugget=!no_nugget, ini_cov_pars=c(10,500))
#gp_fit_scale = likfit_cache(mle_scale,d,covars, nugget=!no_nugget, ini_cov_pars=c(10,500))
#gp_fit_shape = likfit_cache(mle_shape,d,covars, nugget=!no_nugget, ini_cov_pars=c(.01,500))
#
#cov_pars_loc = with(gp_fit_loc, c(cov.pars, nugget))[if(no_nugget) 1:2 else 1:3]
#cov_pars_scale = with(gp_fit_scale, c(cov.pars, nugget))[if(no_nugget) 1:2 else 1:3]
#cov_pars_shape = with(gp_fit_shape, c(cov.pars, nugget))[if(no_nugget) 1:2 else 1:3]
############################################################


###########################################################
############################################################
y_frech = matrix(as.numeric(as.matrix(y_c)),T)
for(i in 1:S_c)
    y_frech[,i] = gev2frech(y_mat[,i],mle_loc[i],mle_scale[i],mle_shape[i])


M0 = fitcopula(y_frech, d[1:S_c,], "gaussian", "whitmat", nugget = 0, smooth = 0.5)
mle_cop_range = M0$param['range']
############################################################
############################################################

#stop()

############################################################
############################################################
# setup and run stan
# data passed to stan
y_mat[is.na(y_mat)] <- -99
gev_spatial_data = list(
    T=T, S=S_c+S_m, P=P, K = n_beta_knots, R=group_size, 
    G=G_c+G_m, G_c=G_c, G_m=G_m, S_c=S_c, S_m=S_m, y=y_mat,
    covars = covars, dist=dist, dist_knots_to_points = dist_knots_to_points,
    mle_cop_range = mle_cop_range
    )

# init function for the init chain
initf_manual = function()list(
    loc = mle_loc,
    scale = mle_scale,
    shape = rep(0.1,length(mle_shape)),
    intercept_loc =   0,
    intercept_scale = 0,
    intercept_shape = 0.1,
    sill_loc =   0.01,
    sill_scale = 0.01,
    sill_shape = 0.01,
    range_loc =   max(dist)/2,
    range_scale = max(dist)/2,
    range_shape = max(dist)/2,
    nugget_loc =   0.01,
    nugget_scale = 0.01,
    nugget_shape = 0.01,
    range_copula = mle_cop_range,
    basis_loc =   matrix(0.01,P,n_beta_knots),
    basis_scale = matrix(0.01,P,n_beta_knots),
    basis_shape = matrix(0.01,P,n_beta_knots),
    range_basis_loc =   matrix(max(dist)/2,P,n_beta_knots),
    range_basis_scale = matrix(max(dist)/2,P,n_beta_knots),
    range_basis_shape = matrix(max(dist)/2,P,n_beta_knots)
    #range_basis_loc = rep(max(dist)/2,P),
    #range_basis_scale = rep(max(dist)/2,P),
    #range_basis_shape = rep(max(dist)/2,P)
    )

if(postprocess_only){
    stanfit = readRDS(posterior_data)
}else{
    stan_model = stan(model_file, data = gev_spatial_data, chains = 0)

    lp_fun <- function(y) {lp = log_prob(stan_model, y); return(ifelse(is.na(lp),-Inf,lp))}
    gr_lp_fun <- function(y) grad_log_prob(stan_model, y)
    lp_grad <- function(y,grad)if(grad) gr_lp_fun(y) else lp_fun(y)

    n_each_par = sapply(stan_model@par_dims,function(x)ifelse(length(x)==0,1,prod(x)))
    n_each_par = n_each_par[-length(n_each_par)]
    n_pars = sum(n_each_par)

    initf = function(){
        lp__ = -Inf
        while(is.infinite(lp__)){
            init = runif(n_pars,-1,1)
            lp__ = lp_fun(init)
        }
        constrain_pars(stan_model,init)
    } 

    #stop()

    # run chains in parallel
    sflist = mclapply(1:nchains, mc.cores = nchains, 
               function(i) stan(fit = stan_model, data = gev_spatial_data, 
                                chains = 1, chain_id = i, iter=iterations, warmup=warmup,
                                seed=rng_seed, refresh=1, init=initf_manual,
                                sample_file = sprintf('%s_%02d',sample_file,i), 
                                diagnostic_file = sprintf('%s_%02d',diagnostic_file,i),
                                control=list(max_treedepth=max_treedepth,adapt_delta=adapt_delta)))

    message('This model is: ',model_name)

    # remove any chains that encountered errors 
    #sflist = sflist[sapply(sflist,function(x)length(slot(x,'sim'))>0)]
    stanfit = if(nchains ==1) sflist[[1]] else sflist2stanfit(sflist)

    print(stanfit)
    saveRDS(stanfit,posterior_data)

    message('Maximum Rhat for each chain:')
    print(sapply(sflist,function(x)max(summary(x)$summary[,'Rhat'])))
}

message('This model is: ',model_name)
############################################################
############################################################
#stop()

#options(warn=2)
###########################
###########################
##  Ugly Postprocess_onlying  ##
###########################
###########################

# posterior and trace plots
if(plot_only){
    preds_grid_pars = readRDS(grid_data)
}else{
    #pars_dt = plot_stan_pars(stanfit, trace_plot_dir, posterior_plot_dir)
    posterior_plots(stanfit,posterior_plot_dir)
    trace_plots(stanfit,trace_plot_dir)
    e = rstan::extract(stanfit)
    pars_list2 = lapply(sample(1:nsamples,min(2000,iterations-warmup)),function(i,e)lapply(e,geti,i),e)

    message('This model is: ',model_name)
    ####################################################
    # krig parameters to grid locations from knots
    cl = makeCluster(nchains)
    #knots_mat = data.matrix(knots)
    d_grid_mat = data.matrix(d_grid)
    message('Conditional Simulation: loc')
    system.time({pred_grid_loc_mat   = parSapply(cl,pars_list2,predict_krig_loc10 ,covars_pred, covars, 
        d, d_grid_mat,dist_knots_to_points, dist_knots_to_grid)})
    message('Conditional Simulation: scale')
    system.time({pred_grid_scale_mat   = parSapply(cl,pars_list2,predict_krig_scale10 ,covars_pred, covars, 
        d, d_grid_mat,dist_knots_to_points, dist_knots_to_grid)})
    message('Conditional Simulation: shape')
    system.time({pred_grid_shape_mat   = parSapply(cl,pars_list2,predict_krig_shape10 ,covars_pred, covars, 
        d, d_grid_mat,dist_knots_to_points, dist_knots_to_grid)})

    pred_grid_loc = as.data.table(melt(pred_grid_loc_mat))
    setnames(pred_grid_loc,c('point','iteration','loc'))
    pred_grid_scale = as.data.table(melt(pred_grid_scale_mat))
    setnames(pred_grid_scale,c('point','iteration','scale'))
    pred_grid_shape = as.data.table(melt(pred_grid_shape_mat))
    setnames(pred_grid_shape,c('point','iteration','shape'))
    ####################################################

    ####################################################
    #stick it all together
    preds_grid_pars = merge(merge(pred_grid_loc,pred_grid_scale,by=c('point','iteration')),pred_grid_shape,by=c('point','iteration'))
    #preds_stns_pars = merge(merge(pred_stns_loc,pred_stns_scale,by=c('point','iteration')),pred_stns_shape,by=c('point','iteration'))

    ####################################################
    # compute return level
    preds_grid_pars[,return_level := loc + scale * ((-log(0.99))^(-shape) - 1)/shape]
    preds_grid_pars[return_level<0, return_level:=0]
    preds_grid_pars[scale<0, scale:=0]
    preds_grid_pars[loc<0, loc:=0]
    saveRDS(preds_grid_pars, grid_data)
}

if(file.exists(quantile_data)){
    mm = readRDS(quantile_data)
}else{
    quantiles = dplyr::summarise(group_by(preds_grid_pars,point),
        q05=quantile(return_level,0.025,na.rm=T),
        q025=quantile(return_level,0.025,na.rm=T),
        q50=quantile(return_level,0.50,na.rm=T),
        q95=quantile(return_level,0.975,na.rm=T),
        q975=quantile(return_level,0.975,na.rm=T),
        q025_loc=quantile(loc,0.025,na.rm=T),
        q50_loc=quantile(loc,0.50,na.rm=T),
        q975_loc=quantile(loc,0.975,na.rm=T),
        q025_scale=quantile(scale,0.025,na.rm=T),
        q50_scale=quantile(scale,0.50,na.rm=T),
        q975_scale=quantile(scale,0.975,na.rm=T),
        q025_shape=quantile(shape,0.025,na.rm=T),
        q50_shape=quantile(shape,0.50,na.rm=T),
        q975_shape=quantile(shape,0.975,na.rm=T),
        q05_loc=quantile(loc,0.025,na.rm=T),
        q95_loc=quantile(loc,0.975,na.rm=T),
        q05_scale=quantile(scale,0.025,na.rm=T),
        q95_scale=quantile(scale,0.975,na.rm=T),
        q05_shape=quantile(shape,0.025,na.rm=T),
        q95_shape=quantile(shape,0.975,na.rm=T))

    quantiles[,ci95:=q975-q025]
    quantiles[,ci:=q95-q05]
    quantiles[,ci_loc:=q975_loc-q025_loc]
    quantiles[,ci_scale:=q975_scale-q025_scale]
    quantiles[,ci_shape:=q975_shape-q025_shape]
    quantiles[,cir:=q50/ci]

    m = cbind(map_region,quantiles)
    nm = names(m)
    nmn = with(expand.grid(c('q05','q025','q50','q95','q975','ci','ci95','cir'),c('','_loc','_scale','_shape')),paste(Var1,Var2,sep=''))
    no_melt_vars = nm[!(nm %in% nmn)]
    mm = na.omit(melt(m,id.vars=no_melt_vars))
    mm[value<0,value:=0]
    mm[,variable:=factor(variable)]

    saveRDS(mm,quantile_data)
}


####################################################

trans = if(season_ == 'summer') 'log' else 'log'
##############################
# Return Level plot
##############################
ylim = extendrange(stns$lat,f=0.01)
xlim = extendrange(stns$lon,f=0.01)
#breaks = if(season_=='summer') waiver() else c(0,50,100,200,500)/10
west = clip_map_polygons('state',xlim,ylim)
p_quantiles = ggplot(mm[variable %in% c('q025','q50','q975')])+
    geom_raster(aes(lon,lat,fill=value))+
    geom_polygon(aes(lon,lat,group=group),color='grey',fill=NA,data=west, alpha=0.5, size=.3)+
    #geom_point(aes(lon,lat),data=ll)+
    scale_fill_gradientn('100 year return\nlevel [cm]',colours=tim.colors(),
        guide=guide_colorbar(label.hjust=0.5,label.vjust=0.5,barwidth=20),
        trans=trans)+
    #geom_point(aes(lon,lat),data=map_region[knot_ind],color='purple')+
    #borders('world','us')+
    coord_quickmap(xlim=xlim,ylim=ylim)+
    facet_grid(~variable)+
    theme_minimal()+
    theme(plot.margin = unit(c(0,0,0, 0),"cm"),legend.position="top")

ggsave(sprintf('%s/return_level_%s.pdf',model_plot_dir,season_),p_quantiles,width=8,height=4)


##############################
# Return Level credible interval
##############################
ylim = extendrange(stns$lat,f=0.01)
xlim = extendrange(stns$lon,f=0.01)
west = clip_map_polygons('state',xlim,ylim)

#breaks = if(season_=='summer') c(1, 5, 20, 150)/10 else c(0,50,100, 200,300,500, 650)/10
q50 = mm[variable %in% c('q50')]
q50[value<1,value:=0]
r = range(range(q50$value))
lr = log(c(floor(r[1]),ceiling(r[2])))
breaks1 = round(exp(seq(lr[1],lr[2],len=6)))

ci = mm[variable %in% c('ci')]
r = range(ci$value)
lr = log(c(max(floor(r[1]),1),ceiling(r[2])))
breaks2 = round(exp(seq(lr[1],lr[2],len=6)))

limits = range(c(breaks1,breaks2))

#clim = c(1,range(q50$value)[2])
p_q50 = ggplot(q50)+
    geom_raster(aes(lon,lat,fill=value))+
    geom_polygon(aes(lon,lat,group=group),color='grey',fill=NA,data=west, alpha=0.5, size=.3)+
    #geom_point(aes(lon,lat),data=ll)+
    scale_fill_gradientn('100 year return\nlevel [cm]',colours=tim.colors(),
        guide=guide_colorbar(title.hjust=0.5,title.vjust=0.5,barwidth=10),
        trans=trans,breaks=breaks2,limits=limits)+
    #geom_point(aes(lon,lat),data=map_region[knot_ind],color='purple')+
    #borders('world','us')+
    coord_quickmap(xlim=xlim,ylim=ylim)+
    #facet_grid(~variable)+
    theme_minimal()+
    theme(plot.margin = unit(c(0,0,0, 0),"cm"),legend.position="top")

#breaks = if(season_=='summer') c(1, 10, 100, 500)/10 else c(1,5,10,25,50,100,300)/10
p_ci = ggplot(ci)+
    geom_raster(aes(lon,lat,fill=value))+
    geom_polygon(aes(lon,lat,group=group),color='grey',fill=NA,data=west, alpha=0.5, size=.3)+
    #geom_point(aes(lon,lat),data=ll)+
    scale_fill_gradientn('100 year return\nlevel 90% CI [cm]',colours=tim.colors(),
        guide=guide_colorbar(title.hjust=0.5,title.vjust=0.5,barwidth=10),
        trans=trans,breaks=breaks2,limits=limits)+
    #geom_point(aes(lon,lat),data=map_region[knot_ind],color='purple')+
    #borders('world','us')+
    coord_quickmap(xlim=xlim,ylim=ylim)+
    #facet_grid(~variable)+
    theme_minimal()+
    theme(legend.position="top",panel.grid=element_blank(),
        plot.margin = unit(c(0,0,0, 0),"cm"))
    #guides(fill=guide_legend(title.vjust=0.5))

pdf(sprintf('%s/return_level_ci_%s.pdf',model_plot_dir,season_),width=8,height=5)
multiplot(p_q50,p_ci,cols=2)
dev.off()

##############################
# Error regions
#############################
cir = mm[variable %in% c('cir')]
r = range(cir$value)
breaks = round(seq(r[1],r[2],len=6))

p_cir = ggplot(cir)+
    geom_raster(aes(lon,lat,fill=value))+
    geom_polygon(aes(lon,lat,group=group),color='grey',fill=NA,data=west, alpha=0.5, size=.3)+
    #geom_point(aes(lon,lat),data=ll)+
    scale_fill_gradientn('100 year return\nlevel 90% CI [cm]',colours=tim.colors(),
        guide=guide_colorbar(title.hjust=0.5,title.vjust=0.5,barwidth=10),
        breaks=breaks,limits=range(breaks))+
    #geom_point(aes(lon,lat),data=map_region[knot_ind],color='purple')+
    #borders('world','us')+
    coord_quickmap(xlim=xlim,ylim=ylim)+
    #facet_grid(~variable)+
    theme_minimal()+
    theme(legend.position="top",panel.grid=element_blank(),
        plot.margin = unit(c(0,0,0, 0),"cm"))
    #guides(fill=guide_legend(title.vjust=0.5))

ggsave(sprintf('%s/ci_ratio_%s.pdf',model_plot_dir,season_),p_cir,width=4,height=5)


##############################
# Parameter only plots
##############################
for(par in c('shape','scale','loc')){
    color = #if(par == 'shape')
            #scale_fill_gradientn(colours=tim.colors(),limits=c(0,0.6))
        #else
            scale_fill_gradientn(colours=tim.colors())
    p_par = ggplot(mm[variable %in% paste(c('q025','q50','q975'),par,sep='_')])+
        geom_raster(aes(lon,lat,fill=value))+
        facet_wrap(~variable)+ 
        color + 
        coord_quickmap(xlim=r_lon,ylim=r_lat)+
        theme_minimal()+
        #theme(legend.position="top",panel.grid=element_blank())+
        xlab('')+ylab('')
    ggsave(sprintf('%s/gev_%s.pdf',model_plot_dir,par),p_par,width=7,height=3)
}

color_breaks = if(season_ == 'summer') waiver() else c(0,10,50,100,300)/10
p_loc_q50 = ggplot(mm[variable %in% c('q50_loc')])+
    geom_raster(aes(lon,lat,fill=value))+
    scale_fill_gradientn('Location',colours=tim.colors(),
        guide=guide_colorbar(title.vjust=0.5,barwidth=10))+
    coord_quickmap(xlim=r_lon,ylim=r_lat)+
    theme_minimal()+
    theme(legend.position="top")
    #ggtitle('Location')

p_scale_q50 = ggplot(mm[variable %in% c('q50_scale')])+
    geom_raster(aes(lon,lat,fill=value))+
    scale_fill_gradientn('Scale',colours=tim.colors(),
        guide=guide_colorbar(title.vjust=0.5,barwidth=10))+
    coord_quickmap(xlim=r_lon,ylim=r_lat)+
    theme_minimal()+
    theme(legend.position="top")
    #ggtitle('Scale')

p_shape_q50 = ggplot(mm[variable %in% c('q50_shape')])+
    geom_raster(aes(lon,lat,fill=value))+
    scale_fill_gradientn('Shape',colours=tim.colors(),
        guide=guide_colorbar(title.vjust=0.5,barwidth=10))+
    coord_quickmap(xlim=r_lon,ylim=r_lat)+
    theme_minimal()+
    theme(legend.position="top")
    #ggtitle('Shape')

pdf(sprintf('%s/gev_pars_q50_%s.pdf',model_plot_dir,season_),width=10,height=5)
multiplot(p_loc_q50,p_scale_q50,p_shape_q50,cols=3)
dev.off()


##############################
# Parameter point plots
##############################
ylim = extendrange(stns$lat,f=0.01)
xlim = extendrange(stns$lon,f=0.01)
west = clip_map_polygons('state',xlim,ylim)
colours = tim.colors(64)
p_loc = ggplot(gev_mle)+
    geom_polygon(aes(lon,lat,group=group),color='grey',fill='white',data=west)+
    geom_point(aes(lon,lat,color=loc))+
    theme_minimal()+
    coord_quickmap(xlim=xlim,ylim=ylim)+
    scale_color_gradientn('Location',colours=colours)+
    theme(legend.position="top",panel.grid=element_blank())+
    xlab('')+ylab('')+theme(plot.margin = unit(c(0,0,0, 0),"cm"))
p_scale = ggplot(gev_mle)+
    geom_polygon(aes(lon,lat,group=group),color='grey',fill='white',data=west)+
    geom_point(aes(lon,lat,color=scale))+
    theme_minimal()+
    coord_quickmap(xlim=xlim,ylim=ylim)+
    scale_color_gradientn('Scale',colours=colours)+
    theme(legend.position="top",panel.grid=element_blank())+
    xlab('')+ylab('')+theme(plot.margin = unit(c(0,0,0, 0),"cm"))
p_shape = ggplot(gev_mle)+
    geom_polygon(aes(lon,lat,group=group),color='grey',fill='white',data=west)+
    geom_point(aes(lon,lat,color=shape))+
    theme_minimal()+
    coord_quickmap(xlim=xlim,ylim=ylim)+
    scale_color_gradientn('Shape',colours=colours)+
    theme(legend.position="top",panel.grid=element_blank())+
    xlab('')+ylab('')+theme(plot.margin = unit(c(0,0,0, 0),"cm"))

pdf(sprintf('%s/mle_gev_pars.pdf',model_plot_dir),width=10,height=5)
multiplot(p_loc,p_scale,p_shape,cols=3)
dev.off()

##############################
# Region map
##############################
complete_ind = apply(y,2,function(x)!any(is.na(x)))
stns[,type:=ifelse(complete_ind,'Complete','Incomplete')]
point_types = stns[,c('lon','lat','type'),with=F]
#point_types[,type:=factor(type)]
#point_types[,type:=factor(type,levels(type)[c(3,1,2)])]
ylim = extendrange(stns$lat,f=0.01)#c(30,50)
xlim = extendrange(stns$lon,f=0.01)#c(-125,-105)
west = clip_map_polygons('state',xlim,ylim)
p_stns = ggplot(point_types)+
    geom_polygon(aes(lon,lat,group=group),color='grey',fill='white',data=west, alpha=0.5, size=.3)+
    geom_point(aes(lon,lat,shape=type,color=type),fill='grey')+
    geom_point(aes(lon,lat),data=map_region[knot_ind],shape=8,color='red')+
    scale_shape_manual('',values=c(19,21))+
    scale_color_manual('',values=c('black','black'))+
    theme_minimal()+
    coord_quickmap(xlim=xlim,ylim=ylim)+
    theme(panel.grid=element_blank())
ggsave(sprintf('%s/stations_with_knots.pdf',model_plot_dir),p_stns,width=7,height=7)