"""
YTArray class.



"""
from __future__ import print_function
#-----------------------------------------------------------------------------
# Copyright (c) 2013, yt Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

import copy
import numpy as np

from distutils.version import LooseVersion
from functools import wraps
from numpy import \
    add, subtract, multiply, divide, logaddexp, logaddexp2, true_divide, \
    floor_divide, negative, power, remainder, mod, absolute, rint, \
    sign, conj, exp, exp2, log, log2, log10, expm1, log1p, sqrt, square, \
    reciprocal, sin, cos, tan, arcsin, arccos, arctan, arctan2, \
    hypot, sinh, cosh, tanh, arcsinh, arccosh, arctanh, deg2rad, rad2deg, \
    bitwise_and, bitwise_or, bitwise_xor, invert, left_shift, right_shift, \
    greater, greater_equal, less, less_equal, not_equal, equal, logical_and, \
    logical_or, logical_xor, logical_not, maximum, minimum, fmax, fmin, \
    isreal, iscomplex, isfinite, isinf, isnan, signbit, copysign, nextafter, \
    modf, ldexp, frexp, fmod, floor, ceil, trunc, fabs, spacing

try:
    # numpy 1.13 or newer
    from numpy import positive, divmod as divmod_, isnat, heaviside
except ImportError:
    positive, divmod_, isnat, heaviside = (None,)*4

from yt.units.unit_object import Unit, UnitParseError
from yt.units.unit_registry import UnitRegistry
from yt.units.dimensions import \
    angle, \
    current_mks, \
    dimensionless, \
    em_dimensions
from yt.utilities.exceptions import \
    YTUnitOperationError, YTUnitConversionError, \
    YTUfuncUnitError, YTIterableUnitCoercionError, \
    YTInvalidUnitEquivalence, YTEquivalentDimsError
from yt.utilities.lru_cache import lru_cache
from numbers import Number as numeric_type
from yt.utilities.on_demand_imports import _astropy
from sympy import Rational
from yt.units.unit_lookup_table import \
    default_unit_symbol_lut
from yt.units.equivalencies import equivalence_registry
from yt.utilities.logger import ytLogger as mylog
from .pint_conversions import convert_pint_units

NULL_UNIT = Unit()
POWER_SIGN_MAPPING = {multiply: 1, divide: -1}

# redefine this here to avoid a circular import from yt.funcs
def iterable(obj):
    try: len(obj)
    except: return False
    return True

def return_arr(func):
    @wraps(func)
    def wrapped(*args, **kwargs):
        ret, units = func(*args, **kwargs)
        if ret.shape == ():
            return YTQuantity(ret, units)
        else:
            # This could be a subclass, so don't call YTArray directly.
            return type(args[0])(ret, units)
    return wrapped

@lru_cache(maxsize=128, typed=False)
def sqrt_unit(unit):
    return unit**0.5

@lru_cache(maxsize=128, typed=False)
def multiply_units(unit1, unit2):
    return unit1 * unit2

def preserve_units(unit1, unit2=None):
    return unit1

@lru_cache(maxsize=128, typed=False)
def power_unit(unit, power):
    return unit**power

@lru_cache(maxsize=128, typed=False)
def square_unit(unit):
    return unit*unit

@lru_cache(maxsize=128, typed=False)
def divide_units(unit1, unit2):
    return unit1/unit2

@lru_cache(maxsize=128, typed=False)
def reciprocal_unit(unit):
    return unit**-1

def passthrough_unit(unit, unit2=None):
    return unit

def return_without_unit(unit, unit2=None):
    return None

def arctan2_unit(unit1, unit2):
    return NULL_UNIT

def comparison_unit(unit1, unit2=None):
    return None

def invert_units(unit):
    raise TypeError(
        "Bit-twiddling operators are not defined for YTArray instances")

def bitop_units(unit1, unit2):
    raise TypeError(
        "Bit-twiddling operators are not defined for YTArray instances")

def get_inp_u_unary(ufunc, inputs, out_arr=None):
    inp = inputs[0]
    u = getattr(inp, 'units', None)
    if u is None:
        u = NULL_UNIT
    if u.dimensions is angle and ufunc in trigonometric_operators:
        inp = inp.in_units('radian').v
        if out_arr is not None:
            out_arr = ufunc(inp).view(np.ndarray)
    return out_arr, inp, u

def get_inp_u_binary(ufunc, inputs):
    inp1 = coerce_iterable_units(inputs[0])
    inp2 = coerce_iterable_units(inputs[1])
    unit1 = getattr(inp1, 'units', None)
    unit2 = getattr(inp2, 'units', None)
    ret_class = get_binary_op_return_class(type(inp1), type(inp2))
    if unit1 is None:
        unit1 = Unit(registry=getattr(unit2, 'registry', None))
    if unit2 is None and ufunc is not power:
        unit2 = Unit(registry=getattr(unit1, 'registry', None))
    elif ufunc is power:
        unit2 = inp2
        if isinstance(unit2, np.ndarray):
            if isinstance(unit2, YTArray):
                if unit2.units.is_dimensionless:
                    pass
                else:
                    raise YTUnitOperationError(ufunc, unit1, unit2)
            unit2 = 1.0
    return (inp1, inp2), (unit1, unit2), ret_class

def handle_preserve_units(inps, units, ufunc, ret_class):
    if units[0] != units[1]:
        any_nonzero = [np.any(inps[0]), np.any(inps[1])]
        if any_nonzero[0] == np.bool_(False):
            units = (units[1], units[1])
        elif any_nonzero[1] == np.bool_(False):
            units = (units[0], units[0])
        else:
            if not units[0].same_dimensions_as(units[1]):
                raise YTUnitOperationError(ufunc, *units)
            inps = (inps[0], ret_class(inps[1]).to(
                ret_class(inps[0]).units))
    return inps, units

def handle_comparison_units(inps, units, ufunc, ret_class, raise_error=False):
    if units[0] != units[1]:
        u1d = units[0].is_dimensionless
        u2d = units[1].is_dimensionless
        any_nonzero = [np.any(inps[0]), np.any(inps[1])]
        if any_nonzero[0] == np.bool_(False):
            units = (units[1], units[1])
        elif any_nonzero[1] == np.bool_(False):
            units = (units[0], units[0])
        elif not any([u1d, u2d]):
            if not units[0].same_dimensions_as(units[1]):
                raise YTUnitOperationError(ufunc, *units)
            else:
                if raise_error:
                    raise YTUfuncUnitError(ufunc, *units)
                inps = (inps[0], ret_class(inps[1]).to(
                    ret_class(inps[0]).units))
    return inps, units

def handle_multiply_divide_units(unit, units, out, out_arr):
    if unit.is_dimensionless and unit.base_value != 1.0:
        if not units[0].is_dimensionless:
            if units[0].dimensions == units[1].dimensions:
                out_arr = np.multiply(out_arr.view(np.ndarray),
                                      unit.base_value, out=out)
                unit = Unit(registry=unit.registry)
    return out, out_arr, unit

def coerce_iterable_units(input_object):
    if isinstance(input_object, np.ndarray):
        return input_object
    if iterable(input_object):
        if any([isinstance(o, YTArray) for o in input_object]):
            ff = getattr(input_object[0], 'units', NULL_UNIT, )
            if any([ff != getattr(_, 'units', NULL_UNIT) for _ in input_object]):
                raise YTIterableUnitCoercionError(input_object)
            # This will create a copy of the data in the iterable.
            return YTArray(input_object)
        return input_object
    else:
        return input_object

def sanitize_units_mul(this_object, other_object):
    inp = coerce_iterable_units(this_object)
    ret = coerce_iterable_units(other_object)
    # If the other object is a YTArray and has the same dimensions as the object
    # under consideration, convert so we don't mix units with the same
    # dimensions.
    if isinstance(ret, YTArray):
        if inp.units.same_dimensions_as(ret.units):
            ret.in_units(inp.units)
    return ret

def sanitize_units_add(this_object, other_object, op_string):
    inp = coerce_iterable_units(this_object)
    ret = coerce_iterable_units(other_object)
    # Make sure the other object is a YTArray before we use the `units`
    # attribute.
    if isinstance(ret, YTArray):
        if not inp.units.same_dimensions_as(ret.units):
            # handle special case of adding or subtracting with zero or
            # array filled with zero
            if not np.any(other_object):
                return ret.view(np.ndarray)
            elif not np.any(this_object):
                return ret
            raise YTUnitOperationError(op_string, inp.units, ret.units)
        ret = ret.in_units(inp.units)
    else:
        # If the other object is not a YTArray, then one of the arrays must be
        # dimensionless or filled with zeros
        if not inp.units.is_dimensionless and np.any(ret):
            raise YTUnitOperationError(op_string, inp.units, dimensionless)
    return ret

def validate_comparison_units(this, other, op_string):
    # Check that other is a YTArray.
    if hasattr(other, 'units'):
        if this.units.expr is other.units.expr:
            if this.units.base_value == other.units.base_value:
                return other
        if not this.units.same_dimensions_as(other.units):
            raise YTUnitOperationError(op_string, this.units, other.units)
        return other.in_units(this.units)

    return other

@lru_cache(maxsize=128, typed=False)
def _unit_repr_check_same(my_units, other_units):
    """
    Takes a Unit object, or string of known unit symbol, and check that it
    is compatible with this quantity. Returns Unit object.

    """
    # let Unit() handle units arg if it's not already a Unit obj.
    if not isinstance(other_units, Unit):
        other_units = Unit(other_units, registry=my_units.registry)

    equiv_dims = em_dimensions.get(my_units.dimensions, None)
    if equiv_dims == other_units.dimensions:
        if current_mks in equiv_dims.free_symbols:
            base = "SI"
        else:
            base = "CGS"
        raise YTEquivalentDimsError(my_units, other_units, base)

    if not my_units.same_dimensions_as(other_units):
        raise YTUnitConversionError(
            my_units, my_units.dimensions, other_units, other_units.dimensions)

    return other_units

unary_operators = (
    negative, absolute, rint, sign, conj, exp, exp2, log, log2,
    log10, expm1, log1p, sqrt, square, reciprocal, sin, cos, tan, arcsin,
    arccos, arctan, sinh, cosh, tanh, arcsinh, arccosh, arctanh, deg2rad,
    rad2deg, invert, logical_not, isreal, iscomplex, isfinite, isinf, isnan,
    signbit, floor, ceil, trunc, modf, frexp, fabs, spacing, positive, isnat,
)

binary_operators = (
    add, subtract, multiply, divide, logaddexp, logaddexp2, true_divide, power,
    remainder, mod, arctan2, hypot, bitwise_and, bitwise_or, bitwise_xor,
    left_shift, right_shift, greater, greater_equal, less, less_equal,
    not_equal, equal, logical_and, logical_or, logical_xor, maximum, minimum,
    fmax, fmin, copysign, nextafter, ldexp, fmod, divmod_, heaviside
)

trigonometric_operators = (
    sin, cos, tan,
)

class YTArray(np.ndarray):
    """
    An ndarray subclass that attaches a symbolic unit object to the array data.

    Parameters
    ----------

    input_array : :obj:`!iterable`
        A tuple, list, or array to attach units to
    input_units : String unit specification, unit symbol object, or astropy units
        The units of the array. Powers must be specified using python
        syntax (cm**3, not cm^3).
    registry : ~yt.units.unit_registry.UnitRegistry
        The registry to create units from. If input_units is already associated
        with a unit registry and this is specified, this will be used instead of
        the registry associated with the unit object.
    dtype : data-type
        The dtype of the array data. Defaults to the dtype of the input data,
        or, if none is found, uses np.float64
    bypass_validation : boolean
        If True, all input validation is skipped. Using this option may produce
        corrupted, invalid units or array data, but can lead to significant
        speedups in the input validation logic adds significant overhead. If set,
        input_units *must* be a valid unit object. Defaults to False.

    Examples
    --------

    >>> from yt import YTArray
    >>> a = YTArray([1, 2, 3], 'cm')
    >>> b = YTArray([4, 5, 6], 'm')
    >>> a + b
    YTArray([ 401.,  502.,  603.]) cm
    >>> b + a
    YTArray([ 4.01,  5.02,  6.03]) m

    NumPy ufuncs will pass through units where appropriate.

    >>> import numpy as np
    >>> a = YTArray(np.arange(8) - 4, 'g/cm**3')
    >>> np.abs(a)
    YTArray([4, 3, 2, 1, 0, 1, 2, 3]) g/cm**3

    and strip them when it would be annoying to deal with them.

    >>> np.log10(a)
    array([       -inf,  0.        ,  0.30103   ,  0.47712125,  0.60205999,
            0.69897   ,  0.77815125,  0.84509804])

    YTArray is tightly integrated with yt datasets:

    >>> import yt
    >>> ds = yt.load('IsolatedGalaxy/galaxy0030/galaxy0030')
    >>> a = ds.arr(np.ones(5), 'code_length')
    >>> a.in_cgs()
    YTArray([  3.08600000e+24,   3.08600000e+24,   3.08600000e+24,
             3.08600000e+24,   3.08600000e+24]) cm

    This is equivalent to:

    >>> b = YTArray(np.ones(5), 'code_length', registry=ds.unit_registry)
    >>> np.all(a == b)
    True

    """
    _ufunc_registry = {
        add: preserve_units,
        subtract: preserve_units,
        multiply: multiply_units,
        divide: divide_units,
        logaddexp: return_without_unit,
        logaddexp2: return_without_unit,
        true_divide: divide_units,
        floor_divide: divide_units,
        negative: passthrough_unit,
        power: power_unit,
        remainder: preserve_units,
        mod: preserve_units,
        fmod: preserve_units,
        absolute: passthrough_unit,
        fabs: passthrough_unit,
        rint: return_without_unit,
        sign: return_without_unit,
        conj: passthrough_unit,
        exp: return_without_unit,
        exp2: return_without_unit,
        log: return_without_unit,
        log2: return_without_unit,
        log10: return_without_unit,
        expm1: return_without_unit,
        log1p: return_without_unit,
        sqrt: sqrt_unit,
        square: square_unit,
        reciprocal: reciprocal_unit,
        sin: return_without_unit,
        cos: return_without_unit,
        tan: return_without_unit,
        sinh: return_without_unit,
        cosh: return_without_unit,
        tanh: return_without_unit,
        arcsin: return_without_unit,
        arccos: return_without_unit,
        arctan: return_without_unit,
        arctan2: arctan2_unit,
        arcsinh: return_without_unit,
        arccosh: return_without_unit,
        arctanh: return_without_unit,
        hypot: preserve_units,
        deg2rad: return_without_unit,
        rad2deg: return_without_unit,
        bitwise_and: bitop_units,
        bitwise_or: bitop_units,
        bitwise_xor: bitop_units,
        invert: invert_units,
        left_shift: bitop_units,
        right_shift: bitop_units,
        greater: comparison_unit,
        greater_equal: comparison_unit,
        less: comparison_unit,
        less_equal: comparison_unit,
        not_equal: comparison_unit,
        equal: comparison_unit,
        logical_and: comparison_unit,
        logical_or: comparison_unit,
        logical_xor: comparison_unit,
        logical_not: return_without_unit,
        maximum: preserve_units,
        minimum: preserve_units,
        fmax: preserve_units,
        fmin: preserve_units,
        isreal: return_without_unit,
        iscomplex: return_without_unit,
        isfinite: return_without_unit,
        isinf: return_without_unit,
        isnan: return_without_unit,
        signbit: return_without_unit,
        copysign: passthrough_unit,
        nextafter: preserve_units,
        modf: passthrough_unit,
        ldexp: bitop_units,
        frexp: return_without_unit,
        floor: passthrough_unit,
        ceil: passthrough_unit,
        trunc: passthrough_unit,
        spacing: passthrough_unit,
        positive: passthrough_unit,
        divmod_: passthrough_unit,
        isnat: return_without_unit,
        heaviside: preserve_units,
    }

    __array_priority__ = 2.0

    def __new__(cls, input_array, input_units=None, registry=None, dtype=None,
                bypass_validation=False):
        if dtype is None:
            dtype = getattr(input_array, 'dtype', np.float64)
        if bypass_validation is True:
            obj = np.asarray(input_array, dtype=dtype).view(cls)
            obj.units = input_units
            if registry is not None:
                obj.units.registry = registry
            return obj
        if input_array is NotImplemented:
            return input_array.view(cls)
        if registry is None and isinstance(input_units, (str, bytes)):
            if input_units.startswith('code_'):
                raise UnitParseError(
                    "Code units used without referring to a dataset. \n"
                    "Perhaps you meant to do something like this instead: \n"
                    "ds.arr(%s, \"%s\")" % (input_array, input_units)
                    )
        if isinstance(input_array, YTArray):
            ret = input_array.view(cls)
            if input_units is None:
                if registry is None:
                    ret.units = input_array.units
                else:
                    units = Unit(str(input_array.units), registry=registry)
                    ret.units = units
            elif isinstance(input_units, Unit):
                ret.units = input_units
            else:
                ret.units = Unit(input_units, registry=registry)
            return ret
        elif isinstance(input_array, np.ndarray):
            pass
        elif iterable(input_array) and input_array:
            if isinstance(input_array[0], YTArray):
                return YTArray(np.array(input_array, dtype=dtype),
                               input_array[0].units, registry=registry)

        # Input array is an already formed ndarray instance
        # We first cast to be our class type

        obj = 
np.asarray(input_array, dtype=dtype)