from __future__ import print_function, division, absolute_import

from .serialize import dask_serialize, dask_deserialize, serialize, deserialize

import netCDF4


@dask_serialize.register(netCDF4.Dataset)
def serialize_netcdf4_dataset(ds):
    # assume mode is read-only
    return {'filename': ds.filepath()}, []


@dask_deserialize.register(netCDF4.Dataset)
def deserialize_netcdf4_dataset(header, frames):
    return netCDF4.Dataset(header['filename'], mode='r')


@dask_serialize.register(netCDF4.Variable)
def serialize_netcdf4_variable(x):
    header, _ = serialize(x.group())
    header['parent-type'] = header['type']
    header['parent-type-serialized'] = header['type-serialized']
    header['name'] = x.name
    return header, []


@dask_deserialize.register(netCDF4.Variable)
def deserialize_netcdf4_variable(header, frames):
    header['type'] = header['parent-type']
    header['type-serialized'] = header['parent-type-serialized']
    parent = deserialize(header, frames)
    return parent.variables[header['name']]


@dask_serialize.register(netCDF4.Group)
def serialize_netcdf4_group(g):
    parent = g
    while parent.parent:
        parent = parent.parent
    header, _ = serialize_netcdf4_dataset(parent)
    header['path'] = g.path
    return header, []


@dask_deserialize.register(netCDF4.Group)
def deserialize_netcdf4_group(header, frames):
    file = deserialize_netcdf4_dataset(header, frames)
    return file[header['path']]