library(tidyverse)
library(lubridate)
library(pbapply)
library(ggmap)

STADIA_MAP_KEY <- "86e80d7a-db93-4bd6-a6f5-eeb11e7c3001"
register_stadiamaps(STADIA_MAP_KEY)

####################
# Q1: Plot the map of India with station markers, and status
# exploratory data analysis
station_hour <- read_csv("./data/raw/station_hour.csv")
station_latlong <- read_csv("./data/raw/stations_latlong.csv")

india_map <- get_stadiamap(
  bbox = c(left = 71, bottom = 8, right = 94, top = 32), # Adjust bounding box as needed
  maptype = "stamen_terrain",
  zoom = 6
)
# Plot the india_map with station markers
ggmap(india_map) +
  geom_point(
    data = station_latlong %>% mutate(Status = ifelse(is.na(Status), "Inactive", Status)),
    aes(x = longitude, y = latitude, fill = Status, shape = Status),
    size = 2, color = "black"
  ) +
  labs(x = "Longitude", y = "Latitude") +
  scale_fill_manual(values = c("Inactive" = "red", "Active" = "green")) +
  scale_shape_manual(values = c("Inactive" = 21, "Active" = 22)) +
  theme_bw() +
  theme(legend.position = "bottom")
ggsave("./figures/india_map.pdf", width = 6, height = 6)


##############
# Q2: Plot the map of Delhi region with station markers

station_delhi <- station_latlong %>% filter(State == "Delhi")
delhi_map <- get_stadiamap(
  bbox = c(left = 76.95, bottom = 28.45, right = 77.4, top = 28.9), # Adjust bounding box as needed
  maptype = "outdoors",
  zoom = 11
)
ggmap(delhi_map) +
  geom_point(
    data = station_delhi,
    aes(x = longitude, y = latitude), size = 3, color = "darkblue"
  ) +
  geom_label(
    data = station_delhi,
    aes(x = longitude, y = latitude, label = StationId),
    size = 3.5, fill = "white", color = "black", vjust = -0.1, hjust = -0.1
  ) +
  labs(x = "Longitude", y = "Latitude") +
  theme_bw()

ggsave("./figures/delhi_map.pdf", width = 8, height = 8)

###################
# Q3: For each station in Delhi region, let's look at the date range and some summary statistics
station_delhi <- station_latlong %>% filter(State == "Delhi")
station_hour %>%
  group_by(StationId) %>%
  summarise(
    count = n(),
    min_date = min(Datetime),
    max_date = max(Datetime),
    missing_pm25 = sum(is.na(PM2.5)) / count,
    missing_pm10 = sum(is.na(PM10)) / count
  ) %>%
  filter(StationId %in% station_delhi$StationId) %>%
  arrange(min_date) %>%
  head(15)

# Conclusion:
#      * The data ranges from around 2015 to 2020 July.
#      * The missing data is about 50% or more for some location. For some locations, PM2.5 is missing for all the data.

####################
# Q4: Let's subset and look at only the data from Delhi region, for the year 2020
df <- station_hour %>%
  filter(Datetime >= "2019-12-01") %>%
  inner_join(station_delhi %>% select(StationId, latitude, longitude), by = "StationId")

nrow(df) # 161,844 observations

df %>% write_csv("./data/station_hour_subset.csv")

#################
# Q5: Let's try to plot the data (since there are 38 stations, we highlight a few representative ones)
df <- read_csv("./data/station_hour_subset.csv")

library(gghighlight)
station_list <- c(
  "DL001", # north delhi, Alipur
  "DL021", # west delhi, Anand Vihar
  "DL002", # east delhi, Dr. Karni Singh Shooting Range
  "DL009", # south delhi, Major Dhyan Chand National Stadium
  "DL018" # central delhi, NSIT Dwarka
)

# for a particular year 2020, create aggregated values for daily average
df %>%
  mutate(Date = date(Datetime)) %>%
  group_by(StationId, Date) %>%
  summarise(PM2.5 = mean(PM2.5, na.rm = T)) %>%
  ggplot() +
  geom_line(aes(x = Date, y = PM2.5, colour = StationId)) +
  scale_y_continuous(trans = "log10", limits = c(7.5, 500)) +
  theme_bw() +
  gghighlight(StationId %in% station_list)

ggsave("./figures/stations_ts.pdf", width = 8, height = 4)


###################
# Q6: Let's look at the lat-long ranges, 
# learn if any processing is needed

df <- read_csv("./data/station_hour_subset.csv")

print(paste("Min lat = ", min(df$latitude), " and Max lat = ", max(df$latitude))) # 28.47 to 28.85
print(paste("Min long = ", min(df$longitude), " and Max long = ", max(df$longitude))) # 76.98 to 77.32

# let's look at some histograms to understand how should we normalized the covariates
hist(df$PM2.5, breaks = 50)
hist(df$NOx, breaks = 50) # A log transformation may be useful as positively skewed

hist(log(1e-5 + df$PM2.5), breaks = 50) # produces good symmetric dist


# let's look at rank correlation structure
M <- df %>% 
  select(-c(StationId, Datetime, AQI, AQI_Bucket, latitude, longitude)) %>%
  drop_na() %>%
  cor(method = "spearman")
corrplot::corrplot(M, type="upper", order="hclust", col = RColorBrewer::brewer.pal(n=8, name="RdYlBu"))

# how about missing data
# let's look at which timepoints, which stations have complete data
covariate_names <- c("PM2.5", "PM10", "NO", "NO2", "NOx", "NH3", "CO", "SO2", "O3", "Benzene", "Toluene", "Xylene")

# create a nice heatmap showcasing the missingness of these data
df %>%
  mutate(Date = date(Datetime)) %>%
  group_by(Date, StationId) %>%
  summarise(
    has_missing = any(is.na(c_across(all_of(covariate_names))))
  ) %>%
  ungroup() %>%
  mutate(
    has_missing = ifelse(has_missing, "Missing", "Complete"),
    Date = as.Date(Date)
  ) %>%
  ggplot(aes(x = StationId, y = Date)) +
  geom_tile(aes(fill = has_missing), color = "white") +
  scale_fill_manual(
    values = c("Complete" = "green", "Missing" = "red"),
    labels = c("COMPLETE", "MISSING") # Capitalize legend labels
  ) +
  theme_bw() +
  labs(x = "StationId", y = "Date", fill = "STATUS") + # Capitalize legend title
  theme(axis.text.x = element_text(angle = 90, hjust = 1))

ggsave("./figures/complete_cases.pdf", width = 8, height = 4)


# What if only PM2.5 and PM10 are considered?
df %>%
  mutate(Date = date(Datetime)) %>%
  group_by(Date, StationId) %>%
  summarise(
    has_missing = any(is.na(c_across(all_of( c("PM2.5", "PM10") ))))
  ) %>%
  ungroup() %>%
  mutate(
    has_missing = ifelse(has_missing, "Missing", "Complete"),
    Date = as.Date(Date)
  ) %>%
  ggplot(aes(x = StationId, y = Date)) +
  geom_tile(aes(fill = has_missing), color = "white") +
  scale_fill_manual(
    values = c("Complete" = "green", "Missing" = "red"),
    labels = c("COMPLETE", "MISSING") # Capitalize legend labels
  ) +
  theme_bw() +
  labs(x = "StationId", y = "Date", fill = "STATUS") +
  theme(axis.text.x = element_text(angle = 90, hjust = 1))

ggsave("./figures/complete_cases-PMvalues.pdf", width = 8, height = 4)


################################
# Q7: Perform the processing
# Based on the exploration, we consider the following
#   - Only use some meaningful variables
#   - get a location (median) and scale (MAD) estimate for each var
#   - Normalize each variable
#   - Create latitude and longtidue variable
#   - Filter the dates from May to June
covariate_names <- c("PM2.5", "PM10", "NH3", "Toluene", "Benzene", "CO", "NO2", "NOx")
loc_med <- df %>% select(any_of(covariate_names)) %>% 
  as.matrix() %>% 
  log1p() %>%
  matrixStats::colMedians(na.rm = T)
scale_mad <- df %>% select(any_of(covariate_names)) %>%
  as.matrix() %>%
  log1p() %>%
  matrixStats::colMads(na.rm = T)

# apply the mutate on each variable
subdf <- df %>% 
  select(any_of(c("StationId", "Datetime", "latitude", "longitude", covariate_names))) %>%
  mutate(
    lat = (latitude - 28.47) / (28.85 - 28.47),
    long = (longitude - 76.98) / (77.32 - 76.98),
    Y = log1p(PM2.5) # log transformation
  ) %>%
  mutate(across(any_of(covariate_names), 
                ~ (log1p(.) -  loc_med[cur_column()]) / scale_mad[cur_column()], 
                .names = "{.col}_trans" ))
subdf %>% write_csv("./data/station_normalized.csv")
