# AUTOGENERATED! DO NOT EDIT! File to edit: 00_tef_core.ipynb (unless otherwise specified).

__all__ = ['logger', 'TEF_object']

# Cell
import math
import numpy as np
import xarray as xr
import logging
from tqdm import tqdm

# Cell
logging.basicConfig(level=logging.INFO,
                    filename="logger.TEF.GETM",
                    filemode="w",
                    format="%(asctime)-15s %(levelname)-8s %(message)s")
logger = logging.getLogger(__name__)

# Cell
class TEF_object():
    """Total Exchange Flow (TEF) Object used for calculating TEF properties"""
    def __init__(self, filename = None, ds = None, **kwargs):

        # What happens if no filename is provided
        if not filename:
            if ds is None:
                self.ds = None
            else:
                self.ds = ds
            return

        # continue with filename
        try:
            self.ds = None
            self.read(filename, **kwargs)
        except (OSError, IOError, RuntimeError):
            try:
                self.read(filename, **kwargs)
            except Exception:
                raise IOError("Unkown file format. Known formats are netcdf.")

        self.transport = None

    def __repr__(self):
        try:
            string = "xarray dataset with {} time steps \n" \
                     "Available fields: {}".format(
                self.timesteps, ", ".join(self.variables)
            )
        except AttributeError:
            string = "Empty TEF object \n" \
                     "Hint: Have you used read() to load data?"
        return string

    def __str__(self):
        return 'Class {}: \n'.format(self.__class__.__name__, self.ds)

    def __len__(self):
        return len(self.ds)

    def __getattr__(self, attr):
        if attr in self.__dict__:
            return getattr(self, attr)
        return getattr(self.ds, attr)

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

    @property
    def timesteps(self):
        if len(self.ds.dims) != 3:
            logger.warning(
                "\nDimensions should be equal to 3, but they are not.\n"
                "You want ... "
            )
            return self.ds.dims[self._get_name_time()]
        return self.ds.dims[self.get_name_time()]

    @property
    def variables(self):
        return list(self.ds.data_vars)

    @property
    def dimensions(self):
        return list(self.ds.dims)

    def read(self, filename,  **kwargs):
        """Reads xarray dataset"""
        if self.ds is None:
            self.ds = xr.open_dataset(filename, **kwargs)
            logger.debug("read: {}".format(self.__str__))
        else:
            raise ValueError("TEF object is already set!")

    def set_up(self,
           time_name=None,
           longitude_name=None,
           latitude_name=None,
           depth_name=None):
        """Setup xarray.dataset and transposes dimensions into needed format"""
        # set dimensions
        if time_name is None:
            self._time_name = self._get_name_time()
        else:
            self._time_name = time_name
        if longitude_name is None:
            self._longitude_name = self._get_name_longitude()
            if self._longitude_name is None:
                self.ds = self.ds.expand_dims("lon")
                self._longitude_name = "lon"
                print("Created dummy dimension for longitude")
        else:
            self._longitude_name = longitude_name

        if latitude_name is None:
            self._latitude_name = self._get_name_latitude()
            if self._latitude_name is None:
                self.ds = self.ds.expand_dims("lat")
                self._latitude_name = "lat"
                print("Created dummy dimension for latitude")
        else:
            self._latitude_name = latitude_name

        if depth_name is None:
            self._depth_name = self._get_name_depth()
            if self._depth_name is None:
                self.ds = self.ds.expand_dims("depth")
                self._depth_name = "depth"
                print("Created dummy dimension for depth")
        else:
            self._depth_name = depth_name
        if time_name is None:
            self._time_name = self._get_name_time()
            if self._time_name is None:
                self.ds = self.ds.expand_dims("time")
                self._time_name = "time"
                print("Created dummy dimension for time")
        else:
            self._time_name = time_name

        # Transpose data
        self.ds = self.ds.transpose(self._time_name,
                                    self._depth_name,
                                    self._latitude_name,
                                    self._longitude_name)

    def _get_name_time(self, x = None):
        """
        check for 'time' dimension and return name
        """
        # check unit
        if x is not None:
            if isinstance(x, np.ndarray):
                logger.info('numpy array -> creating artificial time axis')
                return np.arange(x.shape[0])
        else:
            for dim in self.ds.dims:
                if (('units' in self.ds[dim].attrs and
                    'since' in self.ds[dim].attrs['units']) or
                    ('units' in self.ds[dim].encoding and
                     'since' in self.ds[dim].encoding['units']) or
                    dim in ['time']):
                    return dim
            # check dtype
            for dim in self.ds.variables:
                try:
                    var = self.ds[dim].data[0]
                except IndexError:
                    var = self.ds[dim].data
                if isinstance(var, np.datetime64):
                    return dim
            # no 'time' dimension found
            logger.warning(
                "\n 'time' dimension (dtype='datetime64[ns]') not found."
            )
            return None

    def _get_name_longitude(self):
        """
        check for 'longitude' dimension and return name
        """
        for dim in self.ds.dims:
            if (('units' in self.ds[dim].attrs and
               self.ds[dim].attrs['units'] in ['degree_east', 'degrees_east']) or
               dim in ['lon', 'longitude', 'x']):
                   return dim
        # no 'longitude' dimension found
        logger.warning(
            "\n 'longitude' dimension (unit='degrees_east') not found."
        )
        return None


    def _get_name_latitude(self):
        """
        check for 'latitude' dimension and return name
        """
        for dim in self.ds.dims:
            if (('units' in self.ds[dim].attrs  and
                self.ds[dim].attrs['units'] in ['degree_north', 'degrees_north']) or
                dim in ['lat', 'latitude', 'y']):
                return dim
        # no 'latitude' dimension found
        logger.warning(
            "\n 'latitude' dimension (unit='degrees_north') not found."
        )
        return None

    def _get_name_depth(self):
        """
        check for 'depth' dimension and return name
        """
        for dim in self.ds.dims:
            if (('units' in self.ds[dim].attrs  and
                self.ds[dim].attrs['units'] in ['vertical', 'level']) or
                dim in ['level', 'depth']):
                return dim
        # no 'latitude' dimension found
        logger.warning(
            "\n 'depth not found"
        )
        return None

    from .calc import convert_q_to_Q
    from .calc import sort_1dim
    from .calc import sort_2dim
    from .calc import calc_bulk_values