library("dplyr")
library("tidyr")
library("tidymodels")
library("terra")
library("sf")
library("ecmwfr")
library("ranger")
library("patchwork")

tidymodels_prefer()

# Options
options(scipen=999)
ggplot2::theme_set(ggplot2::theme_bw(base_family="Inter", base_size=16))

train <- read.csv("data/input/train.csv")
train <- train[!names(train) %in% c("X", "Unnamed..0")]
train <- `names<-`(train, tolower(names(train)))

test <- read.csv("data/input/test.csv")
test <- test[!names(test) %in% c("X", "row_id", "cluster")]
test <- `names<-`(test, tolower(names(test)))

prep_features <- function(df) {
    df$year <- as.integer(substr(df$yearmon, 1, 4))
    df$month <- as.integer(substr(df$yearmon, 6, 7))
    df$month_sin <- sin(2 * pi * df$month / 12)
    df$month_cos <- cos(2 * pi * df$month / 12)
    df$loc_id <- paste(df$x, df$y, sep="_")
    
    df
}

train <- prep_features(train)
test <- prep_features(test)

str(train)

str(test)

era_path <- "data/input/era5land_monthly_2001_2024.nc"
dir.create(dirname(era_path), showWarnings=FALSE, recursive=TRUE)

if (!file.exists(era_path)) {
    Sys.setenv(R_KEYRING_BACKEND="env")
    options(keyring_backend="env")
    rc <- readLines(path.expand("~/.cdsapirc"))
    wf_set_key(key=substr(rc[2], 6, 41))
    
    request <- list(
        dataset_short_name="reanalysis-era5-land-monthly-means",
        product_type="monthly_averaged_reanalysis",
        variable=c("total_precipitation", "volumetric_soil_water_layer_1"),
        year=as.character(2001:2024),
        month=sprintf("%02d", 1:12),
        time="00:00",
        area=c(65, -10, 50, 25), # N, W, S, E
        data_format="netcdf",
        download_format="unarchived",
        target=basename(era_path)
    )
    wf_request(request, transfer=TRUE, path=dirname(era_path), time_out=1800, 
               retry=30, verbose=TRUE)
}

sinu <- "+proj=sinu +lon_0=0 +x_0=0 +y_0=0 +a=6371007.181 +b=6371007.181 +units=m +no_defs"

era5 <- rast(era_path)

# Unique stations across train and test, projected to the reanalysis CRS
locs <- bind_rows(
    distinct(train, loc_id, x, y),
    distinct(test, loc_id, x, y)
) |>
    distinct(loc_id, x, y)
locs_v <- vect(locs, geom=c("x", "y"), crs=sinu) |> project("EPSG:4326")

ex <- terra::extract(era5, locs_v, ID=FALSE)
ex$loc_id <- locs$loc_id

# Snap water-cell stations to the nearest land cell (static land mask)
na_idx <- which(is.na(ex[["swvl1_1"]]))
if (length(na_idx) > 0) {
    land_xy <- xyFromCell(era5, which(!is.na(values(era5[[1]], mat=FALSE))))
    na_ll <- crds(locs_v[na_idx])
    snap_xy <- t(vapply(seq_len(nrow(na_ll)), function(i) {
        d <- (land_xy[, 1] - na_ll[i, 1])^2 + (land_xy[, 2] - na_ll[i, 2])^2 # euclidean distance
        land_xy[which.min(d), ]
    }, numeric(2)))
    patch <- terra::extract(era5, snap_xy)
    if ("ID" %in% names(patch)) patch$ID <- NULL
    
    ex[na_idx, names(era5)] <- patch[, names(era5)]
}

# Reshape to (station, month) and build the one-month precipitation lag
months_seq <- format(seq(as.Date("2001-01-01"), 
                         as.Date("2024-12-01"), 
                         by="month"), 
                     "%Y-%m")
era5_cov <- ex |>
    pivot_longer(-loc_id, names_to="layer", values_to="value") |>
    mutate(
        var=sub("_[0-9]+$", "", layer), 
        idx=as.integer(sub("^.*_", "", layer)), 
        yearmon=months_seq[idx]
    ) |>
    select(loc_id, yearmon, var, value) |>
    pivot_wider(names_from=var, values_from=value) |>
    rename(era5_precip=tp, era5_swvl1=swvl1) |>
    arrange(loc_id, yearmon) |>
    group_by(loc_id) |>
    mutate(era5_precip_lag1=dplyr::lag(era5_precip)) |>
    ungroup()

colSums(is.na(era5_cov))

train2 <- train |> left_join(era5_cov, by=c("loc_id", "yearmon"))

train2 |>
    mutate(month=as.integer(substr(yearmon, 6, 7))) |>
    group_by(month) |>
    summarise(`sswc_m (target)`=mean(sswc_m, na.rm=TRUE),
              `ERA5-Land swvl1`=mean(era5_swvl1, na.rm=TRUE),
              .groups="drop") |>
    mutate(across(-month, \(x) as.numeric(scale(x)))) |>
    pivot_longer(-month, names_to="variable", values_to="z") |>
    ggplot(aes(month, z, colour=variable)) +
    geom_line(linewidth=1) +
    geom_point() +
    scale_x_continuous(breaks=1:12) +
    labs(x="Month", y="Standardized monthly mean (z-score)", colour=NULL)

world <- maps::map("world", plot=FALSE, fill=TRUE) |>
    st_as_sf() |>
    st_geometry()

all_locs <- bind_rows(
    train |> distinct(loc_id, x, y),
    test |> distinct(loc_id, x, y)
) |>
    distinct(loc_id, x, y)

loc_sf <- st_as_sf(all_locs, coords=c("x", "y"), crs=sinu)

aoi_sinu <- loc_sf |>
    st_union() |>
    st_convex_hull() |>
    st_buffer(30000) |>
    st_as_sf()

aoi_ext <- ext(vect(aoi_sinu))

bb <- vect(aoi_ext, crs=sinu) |> project("EPSG:4326") |> st_bbox()

dt_sf <- bind_rows(
    mutate(train[c("x", "y")], set="train"),
    mutate(test[c("x", "y")], set="test")
) |>
    st_as_sf(coords=c("x", "y"), crs=sinu)
dt_sf <- dt_sf[!duplicated(dt_sf$geometry), ]
p_space <- ggplot() +
    geom_sf(data=world) +
    geom_sf(data=dt_sf, aes(colour=set), size=0.7, alpha=0.7) +
    scale_colour_manual(values=c(train="steelblue", test="firebrick")) +
    coord_sf(xlim=bb[c(1, 3)], ylim=bb[c(2, 4)]) +
    labs(title="Spatial hold-out design", x=NULL, y=NULL) +
    theme(panel.background=element_rect(fill="lightblue"),
          legend.background=element_rect(fill="white", colour=NA),
          legend.key=element_rect(fill="white", colour=NA))

p_season <- train |>
    group_by(month) |>
    summarise(mean_sswc=mean(sswc_m), .groups="drop") |>
    ggplot(aes(month, mean_sswc)) +
    geom_line(linewidth=1) +
    geom_point() +
    scale_x_continuous(breaks=1:12) +
    labs(title="Seasonal cycle of the response", x="Month", y="Mean sswc_m")

p_space + p_season

rec <- recipe(sswc_m ~ x + y + dlst + dlst.95 + nlst + nlst.95 + geom + evi + 
                  year + month + month_sin + month_cos + era5_swvl1 + 
                  era5_precip + era5_precip_lag1, 
              data=train2) |>
    step_impute_median(dlst, dlst.95, nlst, nlst.95, era5_precip_lag1)

rf_spec <- rand_forest(trees=500) |>
    set_engine("ranger", num.threads=7, seed=123) |>
    set_mode("regression")

wf_rf <- workflow() |> add_recipe(rec) |> add_model(rf_spec)

res_rf <- fit_resamples(wf_rf, resamples=folds, metrics=metrics)
collect_metrics(res_rf)