# Parameters
case_name = "CROCODILE_tutorial_nwa12_MARBL"
CESM_output_dir = (
    "/glade/campaign/cgd/oce/projects/CROCODILE/workshops/2025/Diagnostics/CESM_Output/"
)
start_date = ""
end_date = ""
save_figs = True
fig_output_dir = None
lc_kwargs = {"threads_per_worker": 1}
serial = True
case_root = None
mom6_input_dir = "/glade/campaign/cgd/oce/projects/CROCODILE/workshops/2025/Diagnostics/Input_Dir/CROCODILE_tutorial_nwa12_MARBL_ocnice"
obc_file_str = "forcing_obc_segment_00*.nc"
subset_kwargs = {}
product = "/glade/work/ajanney/crocodile_2025/CUPiD/examples/regional_ocean/computed_notebooks//ocn/Regional_Ocean_OBC.ipynb"

## Apply time boundaries
if len(start_date) > 0 and len(end_date) > 0:
    import cftime

    calendar = sfc_data.time.encoding.get("calendar", "standard")

    calendar_map = {
        "gregorian": cftime.DatetimeProlepticGregorian,
        "noleap": cftime.DatetimeNoLeap,
    }

    CFTime = calendar_map.get(calendar, cftime.DatetimeGregorian)
    y, m, d = [int(i) for i in start_date.split("-")]
    start_date_time = CFTime(y, m, d)
    y, m, d = [int(i) for i in end_date.split("-")]
    end_date_time = CFTime(y, m, d)

    sfc_data = sfc_data.sel(time=slice(start_date_time, end_date_time))

input_dir = None

for file in ["user_nl_mom", "BuildConf/momconf/MOM_input"]:
    filepath = Path(os.path.join(case_root, file))

    if filepath.is_file():
        with open(filepath, "r") as f:
            for line in f:
                if line.strip().startswith("INPUTDIR"):
                    key, value = line.split("=", 1)

                    input_dir = value.strip()
                    break

    if Path(input_dir).is_dir():
        print(f"Found INPUTDIR in {file}")
        break

# Override if autmatic method doesn't work:
# input_dir = user/specified/path

if input_dir is None:
    print(
        "INPUTDIR not found. Try Setting it manually in the `Regional_Ocean_OBC` notebook."
    )
    raise FileNotFoundError(
        "INPUTDIR not found. Try Setting it manually in the `Regional_Ocean_OBC` notebook."
    )

print(f"INPUTDIR: {input_dir}")

boundary_conds = glob.glob(os.path.join(input_dir, "forcing_obc_segment_*.nc"))
print(f"Found {len(boundary_conds)} OBC files:", boundary_conds)

if len(boundary_conds) == 0:
    raise FileNotFoundError(f"No boundary condition files found at {input_dir}")

segment_num_to_dir = {
    "001": "south",
    "002": "north",
    "003": "west",
    "004": "east",
}
segment_dir_to_num = {v: k for k, v in segment_num_to_dir.items()}

list_boundaries = []
temp_time_slice = slice(0, 5)

for boundary_file in boundary_conds:

    boundary = xr.open_mfdataset(boundary_file)

    # Temp time slice
    boundary = boundary.isel(time=temp_time_slice)

    code = (boundary_file.split("_")[-1]).split(".")[0]
    bound = segment_num_to_dir[code]

    # Variables on the U/V grid (corners)
    uv_vars = [f"u_segment_{code}", f"v_segment_{code}"]

    # Variables on the T grid (cell center)
    tracer_vars = [
        f"salt_segment_{code}",
        f"temp_segment_{code}",
        f"eta_segment_{code}",
    ]

    # Create separate datasets for each grid type
    boundary_uv = boundary[uv_vars]
    boundary_t = boundary[tracer_vars]

    # Rename coordinates for both datasets
    boundary_uv = boundary_uv.rename(
        {
            f"ny_segment_{code}": "yq",
            f"nx_segment_{code}": "xq",
            f"lon_segment_{code}": "longitude_q",
            f"lat_segment_{code}": "latitude_q",
        }
    )
    boundary_t = boundary_t.rename(
        {
            f"ny_segment_{code}": "yh",
            f"nx_segment_{code}": "xh",
            f"lon_segment_{code}": "longitude_h",
            f"lat_segment_{code}": "latitude_h",
        }
    )

    # Drop variables and handle squeeze for both
    boundary_uv = boundary_uv.drop_vars(["nxp", "nyp"])
    boundary_t = boundary_t.drop_vars(["nxp", "nyp"])

    if bound in ["north", "south"]:
        # Slicing for a C-grid
        boundary_uv = boundary_uv.isel(xq=slice(0, None, 2))
        boundary_t = boundary_t.isel(xh=slice(1, None, 2))

        # Assign new coordinates
        boundary_uv = boundary_uv.assign_coords(
            xq=("xq", boundary_uv["longitude_q"].values)
        )
        boundary_t = boundary_t.assign_coords(
            xh=("xh", boundary_t["longitude_h"].values)
        )

        # boundary_uv = boundary_uv.squeeze('yq', drop = True)
        # boundary_t = boundary_t.squeeze('yh', drop = True)
    else:
        # Slicing for a C-grid
        boundary_uv = boundary_uv.isel(yq=slice(0, None, 2))
        boundary_t = boundary_t.isel(yh=slice(1, None, 2))

        # Assign new coordinates
        boundary_uv = boundary_uv.assign_coords(
            yq=("yq", boundary_uv["latitude_q"].values)
        )
        boundary_t = boundary_t.assign_coords(
            yh=("yh", boundary_t["latitude_h"].values)
        )

        # boundary_uv = boundary_uv.squeeze('xq', drop = True)
        # boundary_t = boundary_t.squeeze('xh', drop = True)

    # Combine the processed datasets
    boundary = xr.merge([boundary_uv, boundary_t])

    # Remove longitude and latitude variables from the datasets
    for coord in ["longitude_q", "latitude_q"]:
        if coord in boundary:
            boundary = boundary.drop_vars(coord)
    for coord in ["longitude_h", "latitude_h"]:
        if coord in boundary:
            boundary = boundary.drop_vars(coord)

    # Dumb check to make sure OBCs are on the same grid as model output
    model_xh = static_data["xh"]
    model_yh = static_data["yh"]
    model_xq = static_data["xq"]
    model_yq = static_data["yq"]

    # if boundary['xh'].size > 1: assert np.allclose(boundary['xh'].values, model_xh.values, atol=1e-6), f"Boundary {bound} x-coordinates do not match model grid"
    # if boundary['yh'].size > 1: assert np.allclose(boundary['yh'].values, model_yh.values, atol=1e-6), f"Boundary {bound} y-coordinates do not match model grid"
    # if boundary['xq'].size > 1: assert np.allclose(boundary['xq'].values, model_xq.values, atol=1e-6), f"Boundary {bound} corner x-coordinates do not match model grid"
    # if boundary['yq'].size > 1: assert np.allclose(boundary['yq'].values, model_yq.values, atol=1e-6), f"Boundary {bound} corner y-coordinates do not match model grid"

    z_map = [
        f"nz_segment_{code}_u",
        f"nz_segment_{code}_v",
        f"nz_segment_{code}_temp",
        f"nz_segment_{code}_salt",
    ]
    new_boundary = {}
    for var in boundary.data_vars:
        da = boundary[var]
        da.load()

        zdim = [str(d) for d in da.dims if str(d) in z_map]
        if len(zdim) > 0:
            da = da.rename({zdim[0]: "z_l"})
        new_boundary[var.split("_")[0]] = da

    new_boundary = xr.Dataset(new_boundary).expand_dims(boundary=[bound])
    list_boundaries.append(new_boundary)

# boundaries = xr.concat(list_boundaries, dim = 'boundary')

list_boundaries[1]

static_data