# coding: utf-8
# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
    Test the Logarithmic Units and Quantities
"""

from __future__ import (absolute_import, unicode_literals, division,
                        print_function)
from ...extern import six
from ...extern.six.moves import zip

import pickle
import itertools

import pytest
import numpy as np
from numpy.testing.utils import assert_allclose

from ...tests.helper import assert_quantity_allclose
from ... import units as u, constants as c

lu_units = [u.dex, u.mag, u.decibel]

lu_subclasses = [u.DexUnit, u.MagUnit, u.DecibelUnit]

lq_subclasses = [u.Dex, u.Magnitude, u.Decibel]

pu_sample = (u.dimensionless_unscaled, u.m, u.g/u.s**2, u.Jy)


class TestLogUnitCreation(object):

    def test_logarithmic_units(self):
        """Check logarithmic units are set up correctly."""
        assert u.dB.to(u.dex) == 0.1
        assert u.dex.to(u.mag) == -2.5
        assert u.mag.to(u.dB) == -4

    @pytest.mark.parametrize('lu_unit, lu_cls', zip(lu_units, lu_subclasses))
    def test_callable_units(self, lu_unit, lu_cls):
        assert isinstance(lu_unit, u.UnitBase)
        assert callable(lu_unit)
        assert lu_unit._function_unit_class is lu_cls

    @pytest.mark.parametrize('lu_unit', lu_units)
    def test_equality_to_normal_unit_for_dimensionless(self, lu_unit):
        lu = lu_unit()
        assert lu == lu._default_function_unit  # eg, MagUnit() == u.mag
        assert lu._default_function_unit == lu  # and u.mag == MagUnit()

    @pytest.mark.parametrize('lu_unit, physical_unit',
                             itertools.product(lu_units, pu_sample))
    def test_call_units(self, lu_unit, physical_unit):
        """Create a LogUnit subclass using the callable unit and physical unit,
        and do basic check that output is right."""
        lu1 = lu_unit(physical_unit)
        assert lu1.physical_unit == physical_unit
        assert lu1.function_unit == lu1._default_function_unit

    def test_call_invalid_unit(self):
        with pytest.raises(TypeError):
            u.mag([])
        with pytest.raises(ValueError):
            u.mag(u.mag())

    @pytest.mark.parametrize('lu_cls, physical_unit', itertools.product(
        lu_subclasses + [u.LogUnit], pu_sample))
    def test_subclass_creation(self, lu_cls, physical_unit):
        """Create a LogUnit subclass object for given physical unit,
        and do basic check that output is right."""
        lu1 = lu_cls(physical_unit)
        assert lu1.physical_unit == physical_unit
        assert lu1.function_unit == lu1._default_function_unit

        lu2 = lu_cls(physical_unit,
                     function_unit=2*lu1._default_function_unit)
        assert lu2.physical_unit == physical_unit
        assert lu2.function_unit == u.Unit(2*lu2._default_function_unit)

        with pytest.raises(ValueError):
            lu_cls(physical_unit, u.m)


def test_predefined_magnitudes():
    assert_quantity_allclose((-21.1*u.STmag).physical,
                             1.*u.erg/u.cm**2/u.s/u.AA)
    assert_quantity_allclose((-48.6*u.ABmag).physical,
                             1.*u.erg/u.cm**2/u.s/u.Hz)
    assert_quantity_allclose((0*u.M_bol).physical, c.L_bol0)
    assert_quantity_allclose((0*u.m_bol).physical,
                             c.L_bol0/(4.*np.pi*(10.*c.pc)**2))


def test_predefined_reinitialisation():
    assert u.mag('ST') == u.STmag
    assert u.mag('AB') == u.ABmag
    assert u.mag('Bol') == u.M_bol
    assert u.mag('bol') == u.m_bol


def test_predefined_string_roundtrip():
    """Ensure roundtripping; see #5015"""
    with u.magnitude_zero_points.enable():
        assert u.Unit(u.STmag.to_string()) == u.STmag
        assert u.Unit(u.ABmag.to_string()) == u.ABmag
        assert u.Unit(u.M_bol.to_string()) == u.M_bol
        assert u.Unit(u.m_bol.to_string()) == u.m_bol


def test_inequality():
    """Check __ne__ works (regresssion for #5342)."""
    lu1 = u.mag(u.Jy)
    lu2 = u.dex(u.Jy)
    lu3 = u.mag(u.Jy**2)
    lu4 = lu3 - lu1
    assert lu1 != lu2
    assert lu1 != lu3
    assert lu1 == lu4


class TestLogUnitStrings(object):

    def test_str(self):
        """Do some spot checks that str, repr, etc. work as expected."""
        lu1 = u.mag(u.Jy)
        assert str(lu1) == 'mag(Jy)'
        assert repr(lu1) == 'Unit("mag(Jy)")'
        assert lu1.to_string('generic') == 'mag(Jy)'
        with pytest.raises(ValueError):
            lu1.to_string('fits')

        lu2 = u.dex()
        assert str(lu2) == 'dex'
        assert repr(lu2) == 'Unit("dex(1)")'
        assert lu2.to_string() == 'dex(1)'

        lu3 = u.MagUnit(u.Jy, function_unit=2*u.mag)
        assert str(lu3) == '2 mag(Jy)'
        assert repr(lu3) == 'MagUnit("Jy", unit="2 mag")'
        assert lu3.to_string() == '2 mag(Jy)'

        lu4 = u.mag(u.ct)
        assert lu4.to_string('generic') == 'mag(ct)'
        assert lu4.to_string('latex') == ('$\\mathrm{mag}$$\\mathrm{\\left( '
                                          '\\mathrm{ct} \\right)}$')
        assert lu4._repr_latex_() == lu4.to_string('latex')


class TestLogUnitConversion(object):
    @pytest.mark.parametrize('lu_unit, physical_unit',
                             itertools.product(lu_units, pu_sample))
    def test_physical_unit_conversion(self, lu_unit, physical_unit):
        """Check various LogUnit subclasses are equivalent and convertible
        to their non-log counterparts."""
        lu1 = lu_unit(physical_unit)
        assert lu1.is_equivalent(physical_unit)
        assert lu1.to(physical_unit, 0.) == 1.

        assert physical_unit.is_equivalent(lu1)
        assert physical_unit.to(lu1, 1.) == 0.

        pu = u.Unit(8.*physical_unit)
        assert lu1.is_equivalent(physical_unit)
        assert lu1.to(pu, 0.) == 0.125

        assert pu.is_equivalent(lu1)
        assert_allclose(pu.to(lu1, 0.125), 0., atol=1.e-15)

        # Check we round-trip.
        value = np.linspace(0., 10., 6)
        assert_allclose(pu.to(lu1, lu1.to(pu, value)), value, atol=1.e-15)
        # And that we're not just returning True all the time.
        pu2 = u.g
        assert not lu1.is_equivalent(pu2)
        with pytest.raises(u.UnitsError):
            lu1.to(pu2)

        assert not pu2.is_equivalent(lu1)
        with pytest.raises(u.UnitsError):
            pu2.to(lu1)

    @pytest.mark.parametrize('lu_unit', lu_units)
    def test_container_unit_conversion(self, lu_unit):
        """Check that conversion to logarithmic units (u.mag, u.dB, u.dex)
        is only possible when the physical unit is dimensionless."""
        values = np.linspace(0., 10., 6)
        lu1 = lu_unit(u.dimensionless_unscaled)
        assert lu1.is_equivalent(lu1.function_unit)
        assert_allclose(lu1.to(lu1.function_unit, values), values)

        lu2 = lu_unit(u.Jy)
        assert not lu2.is_equivalent(lu2.function_unit)
        with pytest.raises(u.UnitsError):
            lu2.to(lu2.function_unit, values)

    @pytest.mark.parametrize(
        'flu_unit, tlu_unit, physical_unit',
        itertools.product(lu_units, lu_units, pu_sample))
    def test_subclass_conversion(self, flu_unit, tlu_unit, physical_unit):
        """Check various LogUnit subclasses are equivalent and convertible
        to each other if they correspond to equivalent physical units."""
        values = 
np.linspace(0., 10., 6)