#!/usr/bin/env python3
"""
Script to make sure that lon and lat values of Scarpelli Mexico
data files are rounded to 5 decimal places.
"""

import numpy as np
import xarray as xr
import gcpy
from glob import glob

# Get the proper reader
reader = gcpy.util.dataset_reader(
    multi_files=False,
    verbose=False
)


# Keep xarray attributes
with xr.set_options(keep_attrs=True):

    # Loop over each file
    for ifile in glob("../*nc*"):

        print(f"Fixing {ifile}")
        
        # Read data into an xr.Dataset object
        dset = reader(
            ifile,
            drop_variables=gcpy.constants.skip_these_vars
        ).load()

        # Create DataArrays for the lon & lat coords
        darr_lon = dset.coords["lon"]
        darr_lat = dset.coords["lat"]

        # Round lon & lat up to 5 decimal places
        lon = np.round(darr_lon.values, 5)
        lat = np.round(darr_lat.values, 5)

        # Store rounded lon/lat back in the DataArrays
        darr_lon.values = lon
        darr_lat.values = lat

        # Merge the updated DataArrays back into the DataSet
        dset = dset.assign_coords({
            "lon", darr_lon,
            "lat", darr_lat,
        })
        
        # Output file name
        ofile = "fixed." + ifile
        dset.to_netcdf("fixed." + ifile)
        
        # Free memory
        del dset