biome_mask <-
  read_rds(here::here("data/biome_mask.rds"))

map <-
  read_rds(here::here("data/map.rds"))

key_biomes <-
  read_rds(here::here("data/key_biomes.rds"))

super_biomes <-
  read_rds(here::here("data/super_biomes.rds"))

super_biome_mask <-
  read_rds(here::here("data/super_biome_mask.rds"))

labels_breaks <- function(i_name) {
  
  if (i_name == "dco2") {
    i_legend_title <- "ΔpCO<sub>2</sub><br>(µatm)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "dfco2") {
    i_legend_title <- "ΔfCO<sub>2</sub><br>(µatm)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "atm_co2") {
    i_legend_title <- "pCO<sub>2,atm</sub><br>(µatm)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "atm_fco2") {
    i_legend_title <- "fCO<sub>2,atm</sub><br>(µatm)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "sol") {
    i_legend_title <- "K<sub>0</sub><br>(mol m<sup>-3</sup> µatm<sup>-1</sup>)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "kw") {
    i_legend_title <- "k<sub>w</sub><br>(m yr<sup>-1</sup>)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "kw_sol") {
    i_legend_title <- "k<sub>w</sub> K<sub>0</sub><br>(mol yr<sup>-1</sup> m<sup>-2</sup> µatm<sup>-1</sup>)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "spco2") {
    i_legend_title <- "pCO<sub>2,ocean</sub><br>(µatm)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "sfco2") {
    i_legend_title <- "fCO<sub>2,ocean</sub><br>(µatm)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "fgco2") {
    i_legend_title <- "FCO<sub>2</sub><br>(mol m<sup>-2</sup> yr<sup>-1</sup>)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "fgco2_hov") {
    i_legend_title <- "FCO<sub>2</sub><br>(PgC deg<sup>-1</sup> yr<sup>-1</sup>)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "fgco2_int") {
    i_legend_title <- "FCO<sub>2</sub><br>(PgC yr<sup>-1</sup>)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "temperature") {
    i_legend_title <- "SST<br>(°C)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "salinity") {
    i_legend_title <- "SSS"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "chl") {
    i_legend_title <- "lg(Chl-a)<br>(lg(mg m<sup>-3</sup>))"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "mld") {
    i_legend_title <- "lg(MLD)<br>(lg(m))"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "press") {
    i_legend_title <- "pressure<sub>atm</sub><br>(Pa)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "wind") {
    i_legend_title <- "Wind <br>(m sec<sup>-1</sup>)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "SSH") {
    i_legend_title <- "SSH <br>(m)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  if (i_name == "fice") {
    i_legend_title <- "Sea ice <br>(%)"
    # i_breaks <- c(-Inf, seq(0, 80, 10), Inf)
    # i_contour_level <- 50
    # i_contour_level_abs <- 2200
  }
  
  all_labels_breaks <- lst(i_legend_title,
                           # i_breaks,
                           # i_contour_level,
                           # i_contour_level_abs
                           )
  
  return(all_labels_breaks)
  
}


# labels_breaks("fgco2")

x_axis_labels <-
  c(
    "dco2" = labels_breaks("dco2")$i_legend_title,
    "dfco2" = labels_breaks("dfco2")$i_legend_title,
    "atm_co2" = labels_breaks("atm_co2")$i_legend_title,
    "atm_fco2" = labels_breaks("atm_fco2")$i_legend_title,
    "sol" = labels_breaks("sol")$i_legend_title,
    "kw" = labels_breaks("kw")$i_legend_title,
    "kw_sol" = labels_breaks("kw_sol")$i_legend_title,
    "spco2" = labels_breaks("spco2")$i_legend_title,
    "sfco2" = labels_breaks("sfco2")$i_legend_title,
    "fgco2_hov" = labels_breaks("fgco2_hov")$i_legend_title,
    "fgco2_int" = labels_breaks("fgco2_int")$i_legend_title,
    "temperature" = labels_breaks("temperature")$i_legend_title,
    "salinity" = labels_breaks("salinity")$i_legend_title,
    "chl" = labels_breaks("chl")$i_legend_title,
    "mld" = labels_breaks("mld")$i_legend_title,
    "press" = labels_breaks("press")$i_legend_title,
    "wind" = labels_breaks("wind")$i_legend_title,
    "SSH" = labels_breaks("SSH")$i_legend_title,
    "fice" = labels_breaks("fice")$i_legend_title
  )

name_quadratic_fit <- c("atm_co2", "atm_fco2", "spco2", "sfco2")

start_year <- 1990

name_divergent <- c("dco2", "dfco2", "fgco2", "fgco2_hov", "fgco2_int")

pco2_product <-
  pco2_product %>%
  filter(year >= start_year)

pco2_product <-
  full_join(pco2_product,
            biome_mask)

# set all values outside biome mask to NA

pco2_product <-
  pco2_product %>%
  mutate(across(-c(lat, lon, time, area, year, month, biome), 
                ~ if_else(is.na(biome), NA, .)))

# apply coarse grid

pco2_product_coarse <-
  m_grid_horizontal_coarse(pco2_product)

pco2_product_coarse <-
  pco2_product_coarse %>%
  select(-c(lon, lat, time, biome)) %>%
  group_by(year, month, lon_grid, lat_grid) %>%
  summarise(across(-area,
                   ~ weighted.mean(., area))) %>%
  ungroup() %>%
  rename(lon = lon_grid, lat = lat_grid)

pco2_product_coarse <-
  pco2_product_coarse %>%
  pivot_longer(-c(year, month, lon, lat)) %>% 
  drop_na() %>%
  pivot_wider()

# compute annual means

pco2_product_coarse_annual <-
  pco2_product_coarse %>%
  select(-month) %>% 
  group_by(year, lon, lat) %>%
  summarise(across(where(is.numeric),
                   ~ mean(.))) %>%
  ungroup()

pco2_product_coarse_annual <-
  pco2_product_coarse_annual %>% 
  pivot_longer(-c(year, lon, lat))

## compute monthly means

pco2_product_coarse_monthly <-
  pco2_product_coarse %>%
  group_by(year, month, lon, lat) %>%
  summarise(across(where(is.numeric),
                   ~ mean(.))) %>%
  ungroup()

pco2_product_coarse_monthly <-
  pco2_product_coarse_monthly %>% 
  pivot_longer(-c(year, month, lon, lat))

pco2_product_monthly_global <-
  pco2_product %>%
  filter(!is.na(fgco2)) %>% 
  mutate(fgco2_int = fgco2) %>% 
  select(-c(lon, lat, year, month, biome)) %>% 
  group_by(time) %>%
  summarise(across(-c(fgco2_int, area),
                   ~ weighted.mean(., area, na.rm = TRUE)),
            across(fgco2_int,
                   ~ sum(. * area, na.rm = TRUE) * 12.01 * 1e-15)) %>%
  ungroup()

pco2_product_monthly_biome <-
  pco2_product %>%
  filter(!is.na(fgco2)) %>% 
  mutate(fgco2_int = fgco2) %>% 
  select(-c(lon, lat, year, month)) %>% 
  group_by(time, biome) %>%
  summarise(across(-c(fgco2_int, area),
                   ~ weighted.mean(., area, na.rm = TRUE)),
            across(fgco2_int,
                   ~ sum(. * area, na.rm = TRUE) * 12.01 * 1e-15)) %>%
  ungroup()


pco2_product_monthly_biome_super <-
  pco2_product %>%
  filter(!is.na(fgco2)) %>% 
  mutate(fgco2_int = fgco2) %>% 
  mutate(
    biome = case_when(
      str_detect(biome, "NA-") ~ "North Atlantic",
      str_detect(biome, "NP-") ~ "North Pacific",
      str_detect(biome, "SO-") ~ "Southern Ocean",
      TRUE ~ "other"
    )
  ) %>%
  filter(biome != "other") %>%
  select(-c(lon, lat, year, month)) %>%
  group_by(time, biome) %>%
  summarise(across(-c(fgco2_int, area),
                   ~ weighted.mean(., area, na.rm = TRUE)),
            across(fgco2_int,
                   ~ sum(. * area, na.rm = TRUE) * 12.01 * 1e-15)) %>%
  ungroup()

pco2_product_monthly <-
  bind_rows(pco2_product_monthly_global %>%
              mutate(biome = "Global"),
            pco2_product_monthly_biome,
            pco2_product_monthly_biome_super)

rm(
  pco2_product_monthly_global,
  pco2_product_monthly_biome,
  pco2_product_monthly_biome_super
)


pco2_product_monthly <-
  pco2_product_monthly %>% 
  filter(!is.na(biome))

pco2_product_monthly <-
  pco2_product_monthly %>%
  mutate(year = year(time),
         month = month(time),
         .after = time)

pco2_product_monthly <-
  pco2_product_monthly %>%
  pivot_longer(-c(time, year, month, biome))

anomaly_determination <- function(df,...) {
  
  group_by <- quos(...)
  # group_by <- quos(lon, lat)
  # df <- pco2_product_coarse_annual
  
  # Linear regression models
  
  df_lm <-
    df %>%
    filter(year != 2023,
           !(name %in% name_quadratic_fit)) %>%
    nest(data = -c(name, !!!group_by)) %>%
    mutate(
      fit = map(data, ~ lm(value ~ year, data = .x)),
      tidied = map(fit, tidy),
      augmented = map(fit, augment)
    )
  
  
  df_lm_year_anom <-
    full_join(
      df_lm %>%
        unnest(tidied) %>%
        select(name, !!!group_by, term, estimate) %>%
        pivot_wider(names_from = term,
                    values_from = estimate) %>%
        mutate(fit = `(Intercept)` + year * 2023) %>%
        select(name, !!!group_by, fit) %>%
        mutate(year = 2023),
      df %>%
        filter(year == 2023,
               !(name %in% name_quadratic_fit))
    ) %>%
    mutate(resid = value - fit)
  
  
  df_lm <-
    bind_rows(
      df_lm %>%
        unnest(augmented) %>%
        select(name, !!!group_by, year, value, fit = .fitted, resid = .resid),
      df_lm_year_anom
    )
  
  rm(df_lm_year_anom)
  
  # Quadratic regression models
  
  df_quadratic <-
    df %>%
    filter(year != 2023,
           name %in% name_quadratic_fit) %>%
    nest(data = -c(name, !!!group_by)) %>%
    mutate(
      fit = map(data, ~ lm(value ~ year + I(year ^ 2), data = .x)),
      tidied = map(fit, tidy),
      augmented = map(fit, augment)
    )
  
  
  df_quadratic_year_anom <-
    full_join(
      df_quadratic %>%
        unnest(tidied) %>%
        select(name, !!!group_by, term, estimate) %>%
        pivot_wider(names_from = term,
                    values_from = estimate) %>%
        mutate(fit = `(Intercept)` + year * 2023 + `I(year^2)` * 2023 ^ 2) %>%
        select(name, !!!group_by, fit) %>%
        mutate(year = 2023),
      df %>%
        filter(year == 2023,
               name %in% name_quadratic_fit)
    ) %>%
    mutate(resid = value - fit)
  
  
  df_quadratic <-
    bind_rows(
      df_quadratic %>%
        unnest(augmented) %>%
        select(name, !!!group_by, year, value, fit = .fitted, resid = .resid),
      df_quadratic_year_anom
    )
  
  rm(df_quadratic_year_anom)
  
  # Join linear and quadratic regression results
  
  df_regression <-
    bind_rows(df_lm,
              df_quadratic)
  
  df_regression <-
    df_regression %>% 
    arrange(year)
  
  rm(df_lm,
     df_quadratic)
  
  
  return(df_regression)
  
}

anomaly_determination <- function(df,...) {
  
  group_by <- quos(...)
  # group_by <- quos(biome)
  # group_by <- quos(lon, lat)
  # df <- pco2_product_coarse_annual
  
  # Climatoligcal mean
  
  df_mean <-
    df %>%
    filter(year < 2023,
           year >= 2023-5) %>%
    group_by(name, !!!group_by) %>%
    summarize(
      fit = mean(value, na.rm = TRUE)
    ) %>% 
    ungroup()
  
  
  df_mean <-
    full_join(
      df_mean,
      df
    ) %>%
    mutate(resid = value - fit)
  
  return(df_mean)
  
}

pco2_product_hovmoeller_monthly_annual_regression <-
  pco2_product_hovmoeller_monthly_annual %>%
  anomaly_determination(lat) %>% 
  filter(!is.na(resid))

  
pco2_product_hovmoeller_monthly_annual_regression %>%
  # filter(name == "mld") %>% 
  group_split(name) %>%
  # head(1) %>%
  map(
    ~ ggplot(data = .x,
             aes(year, lat, fill = resid)) +
      geom_raster() +
      scale_fill_divergent(name = labels_breaks(.x %>% distinct(name))) +
      theme(legend.title = element_markdown()) +
      coord_cartesian(expand = 0) +
      labs(title = "Annual mean anomalies",
           y = "Latitude") +
      theme(axis.title.x = element_blank())
  )

pco2_product_hovmoeller_monthly_regression <-
  pco2_product_hovmoeller_monthly %>%
  select(-c(decimal)) %>% 
  anomaly_determination(lat, month) %>% 
  filter(!is.na(resid))

  
pco2_product_hovmoeller_monthly_regression <-
  pco2_product_hovmoeller_monthly_regression %>%
  mutate(decimal = year + (month - 1) / 12)
  
pco2_product_hovmoeller_monthly_regression %>%
  group_split(name) %>%
  # head(1) %>%
  map(
    ~ ggplot(data = .x,
             aes(decimal, lat, fill = resid)) +
      geom_raster() +
      scale_fill_divergent(name = labels_breaks(.x %>% distinct(name))) +
      theme(legend.title = element_markdown()) +
      coord_cartesian(expand = 0) +
      labs(title = "Monthly mean anomalies",
           y = "Latitude") +
      theme(axis.title.x = element_blank())
  )

pco2_product_hovmoeller_monthly_regression %>%
  write_csv(paste0("../data/","SOM_FFN","_","2023","_anomaly_hovmoeller_monthly.csv"))