EddyTranslate

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Translate eddy Dataset
"""
import zarr
from netCDF4 import Dataset

from py_eddy_tracker import EddyParser
from py_eddy_tracker.observations.observation import EddiesObservations


def id_parser():
    parser = EddyParser(
        "Eddy Translate, Translate eddies from netcdf to zarr or from zarr to netcdf"
    )
    parser.add_argument("filename_in")
    parser.add_argument("filename_out")
    parser.add_argument("--unraw", action="store_true", help="Load unraw data, use only for netcdf."
    "If unraw is active, netcdf is loaded without apply scalefactor and add_offset.")
    return parser


def is_nc(filename):
    return filename.endswith(".nc")


def get_variable_name(filename):
    if is_nc(filename):
        with Dataset(filename) as h:
            return h.variables.keys()
    else:
        h = zarr.open(filename)
        return list(h.keys())


def get_variable(filename, varname, raw=True):
    if is_nc(filename):
        dataset = EddiesObservations.load_from_netcdf(
            filename, raw_data=raw, include_vars=(varname,)
        )
    else:
        dataset = EddiesObservations.load_from_zarr(filename, include_vars=(varname,))
    return dataset


if __name__ == "__main__":
    args = id_parser().parse_args()

    variables = get_variable_name(args.filename_in)
    if not is_nc(args.filename_out):
        h = zarr.open(args.filename_out, "w")
        for varname in variables:
            get_variable(args.filename_in, varname, raw=not args.unraw).to_zarr(h)
    else:
        with Dataset(args.filename_out, "w") as h:
            for varname in variables:
                get_variable(args.filename_in, varname, raw=not args.unraw).to_netcdf(h)