"""
Reading NetCDF data through xarray.

Three things force the open options below:

- `decode_times=False` is mandatory. None of these files carry a decodable time
  axis: the PEM writes `units = "Planetary year"`, the GCM writes the invalid
  `"days since 0000-00-0 00:00:00"`, and XIOS writes `calendar = "user_defined"`.
  A Mars year is not 365 days, so handing any of them to cftime would produce
  confidently wrong dates rather than an error.
- `concat_characters=False` keeps `controle_descriptor(descriptor, description_size)`
  two-dimensional and typed `|S1`, so the "not numeric" rejection still fires on it.
- `decode_coords='all'` promotes `time_centered` and the `*_bounds` variables to
  coordinates. That is desirable, but it also means `data_vars` alone no longer
  lists every variable, and xarray reports coordinates separately from data
  variables, losing file order. `Source.variables` restores it.
"""

from dataclasses import dataclass, field

import numpy as np
import xarray as xr
from netCDF4 import Dataset as _NC4Dataset

OPEN_KWARGS = dict(
    engine='netcdf4',
    decode_times=False,
    decode_timedelta=False,
    decode_coords='all',
    mask_and_scale=True,
    concat_characters=False,
)

# Chunk only variables bigger than this. Below it, dask's graph overhead costs
# more than the streaming saves, and interactive latency matters.
CHUNK_THRESHOLD_BYTES = 32 * 1024 ** 2

# Target size of one chunk. This must be well below CHUNK_THRESHOLD_BYTES:
# dask's own 'auto' aims at ~128 MiB, which would place every variable in these
# files in a single chunk and stream nothing.
TARGET_CHUNK_BYTES = 8 * 1024 ** 2


def unmask(arr):
    """
    Return a plain float array where masked entries become NaN.
    Accepts masked arrays as well as ordinary arrays.
    """
    if hasattr(arr, 'mask'):
        return np.where(arr.mask, np.nan, np.asarray(arr.data, dtype=float))
    return np.asarray(arr, dtype=float)


def file_variable_order(path):
    """
    Variable names in the order the file declares them.

    xarray splits variables into data_vars and coords and does not preserve the
    on-disk order, but that order is what users see listed in interactive mode.
    """
    nc = _NC4Dataset(path, 'r')
    try:
        return list(nc.variables)
    finally:
        nc.close()


@dataclass
class Source:
    """
    An open NetCDF file: the xarray Dataset plus the file's own variable order.
    """
    path: str
    ds: xr.Dataset
    variables: list = field(default_factory=list)

    def __contains__(self, name):
        return name in self.ds.variables

    def __getitem__(self, name):
        return self.ds[name]

    def dims_and_shape(self, name):
        """
        Dimension names and sizes of one variable, in declaration order.
        """
        da = self.ds[name]
        return list(da.dims), list(da.shape)

    def attr(self, name, key, default=None):
        """
        One attribute of one variable, or `default`.
        """
        return self.ds[name].attrs.get(key, default)

    def close(self):
        self.ds.close()

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        self.close()
        return False


def open_source(path):
    """
    Open a NetCDF file for plotting. Raises the underlying exception on failure.
    """
    ds = xr.open_dataset(path, **OPEN_KWARGS)
    order = file_variable_order(path)
    # Anything xarray synthesized that the file did not declare still belongs in
    # the list, appended after the file's own variables.
    known = set(order)
    order = order + [v for v in ds.variables if v not in known]
    return Source(path=path, ds=ds, variables=order)


def maybe_chunk(da, over=()):
    """
    Wrap a variable in dask when it is large enough to be worth streaming,
    splitting it along the dimensions it is about to be reduced over.

    Chunking along the reduction dimension is the whole point: it is what lets
    an average over 669 time steps be accumulated a few slices at a time instead
    of reading the variable whole. Note that dask's own 'auto' is useless here,
    since its ~128 MiB target would make each of these variables a single chunk.
    """
    if da.chunks is not None or da.nbytes < CHUNK_THRESHOLD_BYTES:
        return da

    candidates = [d for d in over if d in da.dims] or list(da.dims)
    # Split the longest available dimension, so each chunk lands near the target
    dim = max(candidates, key=lambda d: da.sizes[d])
    size = da.sizes[dim]
    bytes_per_index = max(1, da.nbytes // max(1, size))
    per_chunk = max(1, min(size, TARGET_CHUNK_BYTES // bytes_per_index))
    return da.chunk({dim: per_chunk})


def materialize(da):
    """
    Force computation, returning a float64 numpy array.

    Every renderer receives float64 regardless of the file's storage type, which
    is what the netCDF4-based implementation did and what keeps averages from
    accumulating in float32.
    """
    values = da.data
    if hasattr(values, 'compute'):
        values = values.compute()
    return unmask(values)