# use only three basin to assign general basin mask
# ie this is not specific to the MLR fitting
basinmask <- basinmask %>%
  filter(MLR_basins == "2") %>%
  select(lat, lon, basin_AIP)

# set name of model to be subsetted
model_IDs <- c("A", "D")

# for loop across variables
variables <-
  c("spco2",
    "fgco2",
    "talkos",
    "dissicos",
    "po4os",
    "sios",
    "sos",
    "tos")


for (i_model_ID in model_IDs) {
  # i_model_ID <- model_IDs[1]
  
  for (i_variable in variables) {
    # i_variable <- variables[1]
    
    # read list of all files
    file <- paste(i_variable,
                  "_CESM-ETHZ_",
                  i_model_ID,
                  "_1_gr_1980-2018_v20210607.nc",
                  sep = "")
    print(file)
    
    
    # read in data
    variable_data <-
      read_ncdf(paste(path_cmorized,
                      file,
                      sep = ""),
                make_units = FALSE)
    
    # convert to tibble
    variable_data_tibble <- variable_data %>%
      as_tibble()
    
    # remove open link to nc file
    rm(variable_data)
    
    # remove na values
    variable_data_tibble <-
      variable_data_tibble %>%
      drop_na()
    
    # harmonize longitudes
    variable_data_tibble <- variable_data_tibble %>%
      mutate(lon = if_else(lon < 20, lon + 360, lon))
    
    # only consider model grids within basinmask
    variable_data_tibble <-
      inner_join(variable_data_tibble, basinmask) %>%
      select(-basin_AIP)
    
    # mutate variables
    variable_data_tibble <- variable_data_tibble %>%
      mutate(year = year(time_mon)) %>%
      select(-time_mon)
    
    
    # calculate annual average variable
    variable_data_tibble_annual_average <- variable_data_tibble %>%
      fgroup_by(lat, lon, year) %>% {
        add_vars(fgroup_vars(., "unique"),
                 fmean(., keep.group_vars = FALSE))
      }
    
    # calculate climatology
    variable_data_tibble_climatology <-
      variable_data_tibble_annual_average %>%
      fselect(-year) %>%
      fgroup_by(lat, lon) %>% {
        add_vars(fgroup_vars(., "unique"),
                 fmean(., keep.group_vars = FALSE))
      }
    
    # surface map of variable
    print(
    map +
      geom_raster(data = variable_data_tibble_climatology,
                  aes(lon, lat, fill = !!sym(i_variable))) +
      scale_fill_viridis_c(name = i_variable)
    )
    
    
    # write raw data file for GLODAP-based subsetting model variables
    variable_data_tibble_annual_average %>%
      write_csv(file = paste(path_preprocessing,
                             "surface_ocean_",
                             i_model_ID,
                             "/",
                             i_variable,
                             ".csv",
                             sep = ""))
    
    
  }
  
}

# set name of model to be subsetted
model_IDs <- c("A", "B", "C", "D")


file <- paste("area",
              "_CESM-ETHZ_",
              "1_gr.nc",
              sep = "")

area_grid <-
  variable_data <-
  read_ncdf(paste(path_cmorized,
                  file,
                  sep = ""),
            make_units = FALSE) %>% 
  as_tibble()

area_grid <- area_grid %>%
  mutate(lon = if_else(lon < 20, lon + 360, lon))

# for loop across variables
variables <-
  c("spco2",
    "Kw",
    "pco2atm",
    "alpha",
    "fice")


for (i_model_ID in model_IDs) {
  # i_model_ID <- model_IDs[1]
  
  for (i_variable in variables) {
    # i_variable <- variables[5]
    
    # read list of all files
    file <- paste(i_variable,
                  "_CESM-ETHZ_",
                  i_model_ID,
                  "_1_gr_1980-2018_v20210607.nc",
                  sep = "")
    print(file)

    # read in data
    variable_data <-
      read_ncdf(paste(path_cmorized,
                      file,
                      sep = ""),
                make_units = FALSE)
    
    # convert to tibble
    variable_data_tibble <- variable_data %>%
      as_tibble()
    
    # remove open link to nc file
    rm(variable_data)
    
    # remove na values
    variable_data_tibble <-
      variable_data_tibble %>%
      drop_na()
    
    # harmonize longitudes
    variable_data_tibble <- variable_data_tibble %>%
      mutate(lon = if_else(lon < 20, lon + 360, lon))
    
    # only consider model grids within basinmask
    variable_data_tibble <-
      inner_join(variable_data_tibble, basinmask) %>%
      select(-basin_AIP)
    
    if (exists("all_variables")) {
      all_variables <- inner_join(all_variables, variable_data_tibble)
    }
    
    if (!exists("all_variables")) {
      all_variables <- variable_data_tibble
    }
    
    
  }
  
  all_variables <- inner_join(all_variables,
                              area_grid)
  
  mol_to_g <- 12
  P <- 1e-15
  m3_to_kg <- 1030
  sec_to_yr <- 60*60*24*365

  unit_conversion_to_PgCyr <- mol_to_g * P * sec_to_yr / m3_to_kg

  all_variables <- all_variables %>%
    mutate(delta_pco2 = spco2 - pco2atm,
           scale = area * Kw * alpha * (1-fice),
           fgco2 = delta_pco2 * scale)
  
  delta_pco2_monthly <- all_variables %>%
    group_by(time_mon) %>%
    summarise(
      scaling_glob = sum(scale),
      fgco2_glob = sum(fgco2),
      delta_pco2_glob = fgco2_glob / scaling_glob
    ) %>%
    ungroup() %>% 
    mutate(fgco2_glob = fgco2_glob * unit_conversion_to_PgCyr)
  
  delta_pco2_annual <- delta_pco2_monthly %>%
    mutate(year = year(time_mon)) %>%
    group_by(year) %>%
    summarise(
      scaling_glob = mean(scaling_glob),
      fgco2_glob = mean(fgco2_glob),
      delta_pco2_glob = mean(delta_pco2_glob)
    ) %>%
    ungroup()
  
  print(
  ggplot() +
    geom_path(data = delta_pco2_monthly,
              aes(decimal_date(time_mon), delta_pco2_glob, col = "monthly")) +
    geom_path(data = delta_pco2_annual,
              aes(year, delta_pco2_glob, col = "annual")) +
    scale_color_brewer(palette = "Set1", name = "Average") +
    labs(title = paste("model_id:",i_model_ID),
         x = "year")
  )
  
  print(
  ggplot() +
    geom_path(data = delta_pco2_monthly,
              aes(decimal_date(time_mon), fgco2_glob, col = "monthly")) +
    geom_path(data = delta_pco2_annual,
              aes(year, fgco2_glob, col = "annual")) +
    scale_color_brewer(palette = "Set1", name = "Average") +
    labs(title = paste("model_id:",i_model_ID),
         x = "year")
  )
  
  print(
    ggplot() +
      geom_path(
        data = delta_pco2_annual,
        aes(
          year,
          scaling_glob * unit_conversion_to_PgCyr
        )
      )
  )
  
  print(
    ggplot() +
      geom_path(
        data = delta_pco2_annual,
        aes(
          year,
          scaling_glob * delta_pco2_glob * unit_conversion_to_PgCyr,
          col = "scaled"
        )
      ) +
      geom_path(
        data = delta_pco2_annual,
        aes(year, fgco2_glob, col = "integrated")
      ) +
      scale_color_brewer(palette = "Set1", name = "Estimate") +
      scale_y_continuous(name = "Air-sea flux [PgC yr-1]") +
      labs(title = paste("model_id:", i_model_ID),
           x = "year")
  )
  
  # write raw data file for GLODAP-based subsetting model variables
  delta_pco2_annual %>%
    select(-c(scaling_glob, fgco2_glob)) %>% 
    write_csv(
      file = paste(
        path_preprocessing,
        "surface_ocean_disequilibrium/",
        i_model_ID,
        "_annual.csv",
        sep = ""
      )
    )
  
  rm(i_variable, all_variables, delta_pco2_annual, delta_pco2_monthly)
  
}