# In November/December 2024, at IMAU, we conducted a low-resolution quasi-equilibrium 
# hysteresis experiment using the ocean model POP. The experiment was initialized from 
# a statistically equilibrated solution of a present-day control simulation. 
# A quasi-equilibrium approach was employed by adding a slowly varying extra surface 
# freshwater flux (F_H) in the North Atlantic, specifically over the region between 
# latitudes 20°N and 50°N. This additional flux was compensated for over the rest 
# of the domain (excluding marginal seas). The surface freshwater flux was then 
# linearly increased at a rate of 3 × 10⁻⁴ Sv per year until model year 1500, 
# reaching a maximum value of F_H = 0.45 Sv.

# This script generates a file containing the extra surface freshwater flux (F_H). 
# However, the POP model requires this variable to be named SFWF_ANOM, not F_H. 
# The model also expects a climatological input, meaning the file includes 12 identical time steps.

# In this case, SFWF_ANOM is defined for the North Atlantic, spanning the region between 
# latitudes 20°N and 50°N, with compensation applied over the rest of the domain (excluding marginal seas). 
# The total integral of SFWF_ANOM multiplied by the area of the grid cells is 3 × 10⁻⁴ Sv 
# in the specified region (North Atlantic, 20°N-50°N) and 0 Sv over the entire domain due to the compensation.

# The user can modify the region, choose whether or not to apply compensation, and adjust the 
# total flux using the parameters 'mask_area', 'compensation_on', and 'target_hosing'.

# The linear increase mentioned above is handled within the POP model itself by multiplying 
# the extra surface freshwater flux in this file by the relative year in the simulation 
# (relative to the start year).

# Author: Michael Kliphuis

import numpy as np
import xarray as xr
import os
import pandas as pd

############################## USER SECTION ###########################################

# Set the total surface freshwater flux (hosing) in the requested region.
target_hosing = 0.0003 * 1e9  # Target integral of SFWF_ANOM * TAREA in m³/s in the hosing region

# Output file name
output_file = 'SFWF_gx1v6_imau.nc'

# Open a POP dataset to extract grid information and masks. This dataset belongs to a
# quasi-equilibrium hysteresis experiment with the CESM climate model.
dataset = xr.open_dataset(
    '/projects/0/prace_imau/prace_2013081679/cesm1_0_5/b.e10.B1850.f19_g17.qe_hosing.001/OUTPUT/ocn/hist/monthly/'
    'b.e10.B1850.f19_g17.qe_hosing.001.pop.h.0001-01.nc'
)

# Extract relevant variables from the dataset
lon = dataset['TLONG']                # Longitude of grid points (384, 320)
lat = dataset['TLAT']                 # Latitude of grid points (384, 320)
tarea = dataset['TAREA'] / 10000      # Area of grid cells in m² (384, 320)
region_mask = dataset['REGION_MASK']  # Ocean region mask (384, 320)

# Define the target hosing region (20N-50N, region_mask 6 means Atlantic Ocean)
mask_area = (lat >= 20) & (lat <= 50) & (region_mask == 6)

# Toggle for compensating hosing over the rest of the ocean
compensation_on = True

######################### END USER SECTION ###########################################

print("\nCreating the file containing extra surface freshwater flux (SFWF_ANOM)...\n")

# Calculate the total area of the hosing region
area_selection = tarea.where(mask_area, other=0)
total_area_selection = area_selection.sum().item()

# If compensation is enabled, calculate the total area of the compensating region
if compensation_on:
    mask_outside_area = (region_mask > 0) & (~mask_area)  # Exclude marginal seas (region_mask > 0)
    area_outside = tarea.where(mask_outside_area, other=0)
    total_area_outside = area_outside.sum().item()
else:
    mask_outside_area = None
    total_area_outside = 0

# Calculate the remaining area for validation purposes
area_rest = tarea.where(~(mask_area | (mask_outside_area if compensation_on else False)), other=0)
total_area_rest = area_rest.sum().item()

# Validate total area calculations
total_area = total_area_selection + total_area_outside + total_area_rest
total_area_grid = tarea.sum().item()

print(f"Total area of hosing region:                     {total_area_selection / 1e12:.2f} * 10^12 m²")
print(f"Total area of compensation region:              {total_area_outside / 1e12:.2f} * 10^12 m²")
print(f"Total area of remaining grid cells:             {total_area_rest / 1e12:.2f} * 10^12 m²")
print(f"------------------------------------------------------------------------------------")
print(f"Sum of all areas:                               {total_area / 1e12:.2f} * 10^12 m²")
print(f"Total area of the entire grid:                  {total_area_grid / 1e12:.2f} * 10^12 m²")
print("\nPlease verify if the areas match as expected.\n")

# Initialize the SFWF_ANOM array with zeros
sfwf = xr.DataArray(np.zeros(tarea.shape), dims=tarea.dims, coords=tarea.coords)

# Calculate SFWF value for the hosing region
sfwf_value = target_hosing / total_area_selection
sfwf = sfwf.where(~mask_area, other=sfwf_value)

# If compensation is enabled, calculate SFWF for the compensating region
if compensation_on:
    sfwf_value_outside = -target_hosing / total_area_outside
    sfwf = sfwf.where(~mask_outside_area, other=sfwf_value_outside)

# Create the time dimension (12 identical time steps for climatology)
time_dim = pd.date_range('2000-01-01', periods=12, freq='ME')
sfwf_anom = np.full((12, *tarea.shape), sfwf)  # 12 time steps, same values per time step

# Create an xarray DataArray for SFWF_ANOM
sfwf_anom_array = xr.DataArray(
    sfwf_anom, dims=('time', 'nlat', 'nlon'),
    coords={'time': time_dim, 'lat': lat, 'lon': lon}
)

# Add missing_value and _FillValue attributes
missing_value = -1e+34

# Check if the output file already exists and remove it
if os.path.exists(output_file):
    os.remove(output_file)

# Create the output dataset
output_dataset = xr.Dataset({
    'TAREA': dataset['TAREA'],  # Original TAREA variable (in cm²)
    'SFWF_ANOM': sfwf_anom_array,  # New SFWF_ANOM variable
})

# Drop unnecessary coordinates
output_dataset = output_dataset.drop_vars(['ULONG', 'ULAT', 'TLONG', 'TLAT'], errors='ignore')

# Add global attributes to the dataset
output_dataset.SFWF_ANOM.attrs['missing_value'] = missing_value
output_dataset.SFWF_ANOM.attrs['_FillValue'] = missing_value

# Write the dataset to a NetCDF file
output_dataset.to_netcdf(output_file)

print(f"Results have been successfully saved in '{output_file}'.")