library(tidyverse)
library(ggplot2)
library(lubridate)
source("./scripts/estimation.R")

# This is common code
df <- read_csv('./data/station_normalized.csv')

covariate_names <- paste0(c("PM2.5", "PM10", "NH3", "Toluene", "Benzene", "CO", "NO2", "NOx"), "_trans")

stations <- df %>%
  group_by(StationId) %>%
  summarise(latitude = first(latitude), longitude = first(longitude)) # just pick one row from each station
distindex <- compute_distance_orders(stations, location_id_col = "StationId")


##############
# Perform different types of prediction

# Compare models:
# Model 1: Covariate = 1 lag values, all locations
# Model 2: Covariate = 1 lag values, except cur location

compare_models_v1 <- function(
    df,
    pred_location, 
    lagg = 1,
    lambda_bandwidth = 49,
    pred_start = "2020-06-14 00:00:00",
    pred_end = "2020-07-01 00:00:00",
    verbose = FALSE
  ) {
  timepoints <- seq(from = as.POSIXct(pred_start, tz = "UTC"), to = as.POSIXct(pred_end, tz = "UTC"), by = "hour")
  
  # get the true data
  true_df <- df %>%
    filter((Datetime >= timepoints[1]) & (Datetime <= timepoints[length(timepoints)]) & (StationId == pred_location)) %>%
    select(Datetime, Y)
  
  # at time t, covariates from (t-lag) are known
  covdf <- df %>% 
    group_by(StationId) %>%
    mutate(across(any_of(covariate_names), ~ lag(., n = lagg) )) %>%
    ungroup()

  # fit model 1: (with location)
  if (verbose) {
    cat(paste0("Fitting model 1 for location ", pred_location, "\n"))
  }
  mod1 <- model_fit(
    covdf, "Y", timepoints,
    covariate_cols = covariate_names,
    covariate_phi = 0.7,
    lambda_bandwidth = lambda_bandwidth, 
    location_id_column = "StationId",
    verbose = verbose
  )
  
  preddf_mod1 <- model_predict(
    mod1,
    pred_locations = df %>% filter(StationId == pred_location) %>% head(1),
    scale_factor = c(
      quantile(true_df$Y, 0.5, na.rm = TRUE),
      quantile(true_df$Y, 0.75, na.rm = TRUE) - quantile(true_df$Y, 0.25, na.rm = TRUE)
    )
  )
  
  # fit model 2: (without location)
  if (verbose) {
    cat(paste0("Fitting model 2 for location ", pred_location, "\n"))
  }
  mod2 <- model_fit(
    covdf %>% filter(StationId != pred_location), 
    "Y", timepoints,
    covariate_cols = covariate_names,
    covariate_phi = 0.7,
    lambda_bandwidth = lambda_bandwidth, 
    location_id_column = "StationId",
    verbose = verbose
  )
  
  preddf_mod2 <- model_predict(
    mod2,
    pred_locations = df %>% filter(StationId == pred_location) %>% head(1),
    scale_factor = c(
      quantile(true_df$Y, 0.5, na.rm = TRUE),
      quantile(true_df$Y, 0.75, na.rm = TRUE) - quantile(true_df$Y, 0.25, na.rm = TRUE)
    )
  )
  
  # merge the predictions and get residual series
  mergeddf <- preddf_mod1 %>%
    mutate(muY1 = muY) %>%
    select(c("Datetime", "muY1")) %>%
    left_join(preddf_mod2 %>% mutate(muY2 = muY) %>% select(c("Datetime", "muY2")), by = "Datetime") %>%
    left_join(true_df, by = "Datetime") %>%
    mutate(
      e1 = Y - muY1,
      e2 = Y - muY2,
      StationId = pred_location
    ) %>%
    select(c("Datetime", "StationId", "Y", "e1", "e2"))

  return(mergeddf)
}

{
  location_in_out_df <- lapply(unique(stations$StationId), FUN = function(x) {
    result <- compare_models_v1(df, x, 
                lambda_bandwidth = 49, 
                pred_start = "2020-01-01 00:00:00",
                pred_end = "2020-06-30 00:00:00",
                verbose = TRUE
              )
  })
  location_in_out_df %>% bind_rows() %>%
    write_csv("./data/metric_list/location_in_out_residuals.csv")
}


location_in_out_df <- read_csv("./data/metric_list/location_in_out_residuals.csv")




# Compare model 1 and model2 for all locations
# {
#   location_in_out_metric <- lapply(unique(stations$StationId), FUN = function(x) {
#     result <- compare_models_v1(df, x, 
#                              lambda_bandwidth = 49, 
#                              pred_start = "2020-01-01 00:00:00",
#                              pred_end = "2020-06-30 00:00:00",
#                              verbose = TRUE)
#     return(result)
#   })
#   names(location_in_out_metric) <- unique(stations$StationId) # comes in ordered
#   saveRDS(location_in_out_metric, "./data/metric_list/location_in_out_metric.Rds")
# }

# location_in_out_metric <- readRDS("./version3/location_in_out_metric.Rds")







# Take gapped series
{
  take_every <- 1
  skip_every <- 0
  location_in_out_gapped <- lapply(names(location_in_out_metric), function(loc) {
    result <- location_in_out_metric[[loc]]
    if (!all(is.na(result$accuracy)) & take_every > 1 & skip_every > 0) {
      sube1 <- create_gapped_vector(result$e1, take_every, skip_every)
      sube2 <- create_gapped_vector(result$e2, take_every, skip_every)
      subY <- create_gapped_vector(result$Y, take_every, skip_every)
      result2 <- compare_prediction_errors(tibble("e1" = sube1, "e2" = sube2, "Y" = subY))
      return(result2)
    } else {
      return(result)
    }
  })
  
  names(location_in_out_gapped) <- names(location_in_out_metric)
}

table_data <- lapply(names(location_in_out_gapped), function(loc) {
  acc <- location_in_out_gapped[[loc]]$accuracy
  test <- location_in_out_gapped[[loc]]$test
  
  if (!any(is.null(acc) | is.na(acc))) {
    tibble(
      Location = loc,
      MAPE_in  = round(acc["MAPE", "model1"], 3),
      MAPE_out = round(acc["MAPE", "model2"], 3),
      DM_stat  = round(test["dm", "statistic"], 3),
      DM_pval  = round(test["dm", "pvalue"], 3),
      CW_stat  = round(test["cw", "statistic"], 3),
      CW_pval  = round(test["cw", "pvalue"], 3)
    )
  } else {
    tibble(Location = loc, MAPE_in = NA, MAPE_out = NA, DM_stat = NA, DM_pval = NA, CW_stat = NA, CW_pval = NA)
  }
}) %>%
  bind_rows() %>%
  drop_na() %>%
  View()
  # mutate(
  #   MAPE_in  = sprintf("%.3f", MAPE_in * 100),
  #   MAPE_out = sprintf("%.3f", MAPE_out * 100),
  #   DM_comb  = sprintf("%.3f (%.3f)", DM_stat, DM_pval),
  #   CW_comb  = sprintf("%.3f (%.3f)", CW_stat, CW_pval)
  # ) %>%
  # select(Location, MAPE_in, MAPE_out, DM_comb, CW_comb)

# Output LaTeX table code (using kable)
knitr::kable(
  table_data,
  format = "latex",
  booktabs = TRUE,
  col.names = c("Location", "MAPE (In)", "MAPE (Out)", "DM Stat (p-val)", "CW Stat (p-val)")
)


##############
# Do nice prediction based heatmap visualization

head(df)
timepoints <- as.POSIXct(c("2020-01-01 10:00:00", "2020-01-01 20:00:00", "2020-03-01 10:00:00", "2020-03-01 20:00:00", "2020-06-01 10:00:00", "2020-06-01 20:00:00"))

covdf <- df %>% 
  group_by(StationId) %>%
  mutate(across(any_of(covariate_names), ~ lag(., n = 1) )) %>%
  ungroup()

mod <- model_fit(
    covdf, "Y", timepoints,
    covariate_cols = covariate_names,
    covariate_phi = 0.7,
    lambda_bandwidth = 49, 
    location_id_column = "StationId",
    verbose = TRUE
  )

# Define Delhi bounding box (same as your map)
lon_seq <- seq(76.99, 77.31, by = 0.005)   # smaller step -> finer resolution
lat_seq <- seq(28.48, 28.84, by = 0.005)
pred_grid <- expand_grid(longitude = lon_seq, latitude = lat_seq) %>%  # prediction grid
  mutate(
    StationId = paste0("grid_", row_number()),   # give some random stationid
    lat = (latitude - 28.47) / (28.85 - 28.47),   # transformed lat long
    long = (longitude - 76.98) / (77.32 - 76.98)
  )   

preddf <- model_predict(mod, pred_locations = pred_grid)
preddf <- preddf %>% left_join(pred_grid, by = "StationId")
preddf <- preddf %>% 
  mutate(
    tod = if_else(hour(Datetime) < 12, "Morning", "Evening"),  # Define day vs night    
    month_label = month(Datetime, label = TRUE, abbr = FALSE) # Label for the facet header, like "January"
  )

# delhi stations
station_latlong <- read_csv("./data/raw/stations_latlong.csv")
station_delhi <- station_latlong %>% filter(State == "Delhi")

delhi_map <- get_stadiamap(
  bbox = c(left = min(lon_seq), bottom = min(lat_seq), right = max(lon_seq), top = max(lat_seq)), # Adjust bounding box as needed
  maptype = "outdoors",
  zoom = 11
)

ggmap(delhi_map) +
  geom_tile(data = preddf, aes(x = longitude, y = latitude, fill = muY), alpha = 0.5) +
  # MetBrewer::scale_fill_met_c("Veronese", direction = 1, name = "Pollution") +
  scale_fill_viridis_c(option = "plasma", name = "log(PM2.5)\nConcentration") +
  geom_point(
      data = station_delhi,
      aes(x = longitude, y = latitude), size = 1, color = "black", shape = 3
  ) +
  facet_grid(rows = vars(tod), cols = vars(month_label), labeller = label_wrap_gen()) +
  labs(x = "Longitude", y = "Latitude") +
  theme_minimal()

ggsave("./figures/predicted_pm10.pdf", width = 10, height = 6)
