# -*- coding: utf-8 -*-
# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Quantity helpers for the ERFA ufuncs."""
# Tests for these are in coordinates, not in units.

from erfa import dt_eraASTROM, dt_eraLDBODY, dt_pv
from erfa import ufunc as erfa_ufunc

from astropy.units.core import UnitsError, UnitTypeError, dimensionless_unscaled
from astropy.units.structured import StructuredUnit

from . import UFUNC_HELPERS
from .helpers import (
    _d, get_converter, helper_invariant, helper_multiplication, helper_twoarg_invariant)

erfa_ufuncs = ('s2c', 's2p', 'c2s', 'p2s', 'pm', 'pdp', 'pxp', 'rxp',
               'cpv', 'p2pv', 'pv2p', 'pv2s', 'pvdpv', 'pvm', 'pvmpv', 'pvppv',
               'pvstar', 'pvtob', 'pvu', 'pvup', 'pvxpv', 'rxpv', 's2pv', 's2xpv',
               'starpv', 'sxpv', 'trxpv', 'gd2gc', 'gc2gd', 'ldn', 'aper',
               'apio', 'atciq', 'atciqn', 'atciqz', 'aticq', 'atioq', 'atoiq')


def has_matching_structure(unit, dtype):
    dtype_fields = dtype.fields
    if dtype_fields:
        return (isinstance(unit, StructuredUnit)
                and len(unit) == len(dtype_fields)
                and all(has_matching_structure(u, df_v[0])
                        for (u, df_v) in zip(unit.values(), dtype_fields.values())))
    else:
        return not isinstance(unit, StructuredUnit)


def check_structured_unit(unit, dtype):
    if not has_matching_structure(unit, dtype):
        msg = {dt_pv: 'pv',
               dt_eraLDBODY: 'ldbody',
               dt_eraASTROM: 'astrom'}.get(dtype, 'function')
        raise UnitTypeError(f'{msg} input needs unit matching dtype={dtype}.')


def helper_s2c(f, unit1, unit2):
    from astropy.units.si import radian
    try:
        return [get_converter(unit1, radian),
                get_converter(unit2, radian)], dimensionless_unscaled
    except UnitsError:
        raise UnitTypeError("Can only apply '{}' function to "
                            "quantities with angle units"
                            .format(f.__name__))


def helper_s2p(f, unit1, unit2, unit3):
    from astropy.units.si import radian
    try:
        return [get_converter(unit1, radian),
                get_converter(unit2, radian), None], unit3
    except UnitsError:
        raise UnitTypeError("Can only apply '{}' function to "
                            "quantities with angle units"
                            .format(f.__name__))


def helper_c2s(f, unit1):
    from astropy.units.si import radian
    return [None], (radian, radian)


def helper_p2s(f, unit1):
    from astropy.units.si import radian
    return [None], (radian, radian, unit1)


def helper_gc2gd(f, nounit, unit1):
    from astropy.units.si import m, radian
    if nounit is not None:
        raise UnitTypeError("ellipsoid cannot be a quantity.")
    try:
        return [None, get_converter(unit1, m)], (radian, radian, m, None)
    except UnitsError:
        raise UnitTypeError("Can only apply '{}' function to "
                            "quantities with length units"
                            .format(f.__name__))


def helper_gd2gc(f, nounit, unit1, unit2, unit3):
    from astropy.units.si import m, radian
    if nounit is not None:
        raise UnitTypeError("ellipsoid cannot be a quantity.")
    try:
        return [None,
                get_converter(unit1, radian),
                get_converter(unit2, radian),
                get_converter(unit3, m)], (m, None)
    except UnitsError:
        raise UnitTypeError("Can only apply '{}' function to lon, lat "
                            "with angle and height with length units"
                            .format(f.__name__))


def helper_p2pv(f, unit1):
    from astropy.units.si import s
    if isinstance(unit1, StructuredUnit):
        raise UnitTypeError("p vector unit cannot be a structured unit.")
    return [None], StructuredUnit((unit1, unit1 / s))


def helper_pv2p(f, unit1):
    check_structured_unit(unit1, dt_pv)
    return [None], unit1[0]


def helper_pv2s(f, unit_pv):
    from astropy.units.si import radian
    check_structured_unit(unit_pv, dt_pv)
    ang_unit = radian * unit_pv[1] / unit_pv[0]
    return [None], (radian, radian, unit_pv[0], ang_unit, ang_unit, unit_pv[1])


def helper_s2pv(f, unit_theta, unit_phi, unit_r, unit_td, unit_pd, unit_rd):
    from astropy.units.si import radian
    time_unit = unit_r / unit_rd
    return [get_converter(unit_theta, radian),
            get_converter(unit_phi, radian),
            None,
            get_converter(unit_td, radian / time_unit),
            get_converter(unit_pd, radian / time_unit),
            None], StructuredUnit((unit_r, unit_rd))


def helper_pv_multiplication(f, unit1, unit2):
    check_structured_unit(unit1, dt_pv)
    check_structured_unit(unit2, dt_pv)
    result_unit = StructuredUnit((unit1[0] * unit2[0], unit1[1] * unit2[0]))
    converter = get_converter(unit2, StructuredUnit(
        (unit2[0], unit1[1] * unit2[0] / unit1[0])))
    return [None, converter], result_unit


def helper_pvm(f, unit1):
    check_structured_unit(unit1, dt_pv)
    return [None], (unit1[0], unit1[1])


def helper_pvstar(f, unit1):
    from astropy.units.astrophys import AU
    from astropy.units.si import arcsec, day, km, radian, s, year

    return [get_converter(unit1, StructuredUnit((AU, AU/day)))], (
        radian, radian, radian / year, radian / year, arcsec, km / s, None)


def helper_starpv(f, unit_ra, unit_dec, unit_pmr, unit_pmd,
                  unit_px, unit_rv):
    from astropy.units.astrophys import AU
    from astropy.units.si import arcsec, day, km, radian, s, year

    return [get_converter(unit_ra, radian),
            get_converter(unit_dec, radian),
            get_converter(unit_pmr, radian/year),
            get_converter(unit_pmd, radian/year),
            get_converter(unit_px, arcsec),
            get_converter(unit_rv, km/s)], (StructuredUnit((AU, AU/day)), None)


def helper_pvtob(f, unit_elong, unit_phi, unit_hm,
                 unit_xp, unit_yp, unit_sp, unit_theta):
    from astropy.units.si import m, radian, s

    return [get_converter(unit_elong, radian),
            get_converter(unit_phi, radian),
            get_converter(unit_hm, m),
            get_converter(unit_xp, radian),
            get_converter(unit_yp, radian),
            get_converter(unit_sp, radian),
            get_converter(unit_theta, radian)], StructuredUnit((m, m/s))


def helper_pvu(f, unit_t, unit_pv):
    check_structured_unit(unit_pv, dt_pv)
    return [get_converter(unit_t, unit_pv[0]/unit_pv[1]), None], unit_pv


def helper_pvup(f, unit_t, unit_pv):
    check_structured_unit(unit_pv, dt_pv)
    return [get_converter(unit_t, unit_pv[0]/unit_pv[1]), None], unit_pv[0]


def helper_s2xpv(f, unit1, unit2, unit_pv):
    check_structured_unit(unit_pv, dt_pv)
    return [None, None, None], StructuredUnit((_d(unit1) * unit_pv[0],
                                               _d(unit2) * unit_pv[1]))


def ldbody_unit():
    from astropy.units.astrophys import AU, Msun
    from astropy.units.si import day, radian

    return StructuredUnit((Msun, radian, (AU, AU/day)),
                          erfa_ufunc.dt_eraLDBODY)


def astrom_unit():
    from astropy.units.astrophys import AU
    from astropy.units.si import rad, year
    one = rel2c = dimensionless_unscaled

    return StructuredUnit((year, AU, one, AU, rel2c, one, one, rad, rad, rad, rad,
                           one, one, rel2c, rad, rad, rad),
                          erfa_ufunc.dt_eraASTROM)


def helper_ldn(f, unit_b, unit_ob, unit_sc):
    from astropy.units.astrophys import AU
    return [get_converter(unit_b, ldbody_unit()),
            get_converter(unit_ob, AU),
            get_converter(_d(unit_sc), dimensionless_unscaled)], dimensionless_unscaled


def helper_aper(f, unit_theta, unit_astrom):
    check_structured_unit(unit_astrom, dt_eraASTROM)
    unit_along = unit_astrom[7]  # along

    if unit_astrom[14] is unit_along:  # eral
        result_unit = unit_astrom
    else:
        result_units = tuple((unit_along if i == 14 else v)
                             for i, v in enumerate(unit_astrom.values()))
        result_unit = unit_astrom.__class__(result_units, names=unit_astrom)
    return [get_converter(unit_theta, unit_along), None], result_unit


def helper_apio(f, unit_sp, unit_theta, unit_elong, unit_phi, unit_hm,
                unit_xp, unit_yp, unit_refa, unit_refb):
    from astropy.units.si import m, radian
    return [get_converter(unit_sp, radian),
            get_converter(unit_theta, radian),
            get_converter(unit_elong, radian),
            get_converter(unit_phi, radian),
            get_converter(unit_hm, m),
            get_converter(unit_xp, radian),
            get_converter(unit_xp, radian),
            get_converter(unit_xp, radian),
            get_converter(unit_xp, radian)], astrom_unit()


def helper_atciq(f, unit_rc, unit_dc, unit_pr, unit_pd, unit_px, unit_rv, unit_astrom):
    from astropy.units.si import arcsec, km, radian, s, year
    return [get_converter(unit_rc, radian),
            get_converter(unit_dc, radian),
            get_converter(unit_pr, radian / year),
            get_converter(unit_pd, radian / year),
            get_converter(unit_px, arcsec),
            get_converter(unit_rv, km / s),
            get_converter(unit_astrom, astrom_unit())], (radian, radian)


def helper_atciqn(f, unit_rc, unit_dc, unit_pr, unit_pd, unit_px, unit_rv, unit_astrom,
                  unit_b):
    from astropy.units.si import arcsec, km, radian, s, year
    return [get_converter(unit_rc, radian),
            get_converter(unit_dc, radian),
            get_converter(unit_pr, radian / year),
            get_converter(unit_pd, radian / year),
            get_converter(unit_px, arcsec),
            get_converter(unit_rv, km / s),
            get_converter(unit_astrom, astrom_unit()),
            get_converter(unit_b, ldbody_unit())], (radian, radian)


def helper_atciqz_aticq(f, unit_rc, unit_dc, unit_astrom):
    from astropy.units.si import radian
    return [get_converter(unit_rc, radian),
            get_converter(unit_dc, radian),
            get_converter(unit_astrom, astrom_unit())], (radian, radian)


def helper_aticqn(f, unit_rc, unit_dc, unit_astrom, unit_b):
    from astropy.units.si import radian
    return [get_converter(unit_rc, radian),
            get_converter(unit_dc, radian),
            get_converter(unit_astrom, astrom_unit()),
            get_converter(unit_b, ldbody_unit())], (radian, radian)


def helper_atioq(f, unit_rc, unit_dc, unit_astrom):
    from astropy.units.si import radian
    return [get_converter(unit_rc, radian),
            get_converter(unit_dc, radian),
            get_converter(unit_astrom, astrom_unit())], (radian,)*5


def helper_atoiq(f, unit_type, unit_ri, unit_di, unit_astrom):
    from astropy.units.si import radian
    if unit_type is not None:
        raise UnitTypeError("argument 'type' should not have a unit")

    return [None,
            get_converter(unit_ri, radian),
            get_converter(unit_di, radian),
            get_converter(unit_astrom, astrom_unit())], (radian, radian)


def get_erfa_helpers():
    ERFA_HELPERS = {}
    ERFA_HELPERS[erfa_ufunc.s2c] = helper_s2c
    ERFA_HELPERS[erfa_ufunc.s2p] = helper_s2p
    ERFA_HELPERS[erfa_ufunc.c2s] = helper_c2s
    ERFA_HELPERS[erfa_ufunc.p2s] = helper_p2s
    ERFA_HELPERS[erfa_ufunc.pm] = helper_invariant
    ERFA_HELPERS[erfa_ufunc.cpv] = helper_invariant
    ERFA_HELPERS[erfa_ufunc.p2pv] = helper_p2pv
    ERFA_HELPERS[erfa_ufunc.pv2p] = helper_pv2p
    ERFA_HELPERS[erfa_ufunc.pv2s] = helper_pv2s
    ERFA_HELPERS[erfa_ufunc.pvdpv] = helper_pv_multiplication
    ERFA_HELPERS[erfa_ufunc.pvxpv] = helper_pv_multiplication
    ERFA_HELPERS[erfa_ufunc.pvm] = helper_pvm
    ERFA_HELPERS[erfa_ufunc.pvmpv] = helper_twoarg_invariant
    ERFA_HELPERS[erfa_ufunc.pvppv] = helper_twoarg_invariant
    ERFA_HELPERS[erfa_ufunc.pvstar] = helper_pvstar
    ERFA_HELPERS[erfa_ufunc.pvtob] = helper_pvtob
    ERFA_HELPERS[erfa_ufunc.pvu] = helper_pvu
    ERFA_HELPERS[erfa_ufunc.pvup] = helper_pvup
    ERFA_HELPERS[erfa_ufunc.pdp] = helper_multiplication
    ERFA_HELPERS[erfa_ufunc.pxp] = helper_multiplication
    ERFA_HELPERS[erfa_ufunc.rxp] = helper_multiplication
    ERFA_HELPERS[erfa_ufunc.rxpv] = helper_multiplication
    ERFA_HELPERS[erfa_ufunc.s2pv] = helper_s2pv
    ERFA_HELPERS[erfa_ufunc.s2xpv] = helper_s2xpv
    ERFA_HELPERS[erfa_ufunc.starpv] = helper_starpv
    ERFA_HELPERS[erfa_ufunc.sxpv] = helper_multiplication
    ERFA_HELPERS[erfa_ufunc.trxpv] = helper_multiplication
    ERFA_HELPERS[erfa_ufunc.gc2gd] = helper_gc2gd
    ERFA_HELPERS[erfa_ufunc.gd2gc] = helper_gd2gc
    ERFA_HELPERS[erfa_ufunc.ldn] = helper_ldn
    ERFA_HELPERS[erfa_ufunc.aper] = helper_aper
    ERFA_HELPERS[erfa_ufunc.apio] = helper_apio
    ERFA_HELPERS[erfa_ufunc.atciq] = helper_atciq
    ERFA_HELPERS[erfa_ufunc.atciqn] = helper_atciqn
    ERFA_HELPERS[erfa_ufunc.atciqz] = helper_atciqz_aticq
    ERFA_HELPERS[erfa_ufunc.aticq] = helper_atciqz_aticq
    ERFA_HELPERS[erfa_ufunc.aticqn] = helper_aticqn
    ERFA_HELPERS[erfa_ufunc.atioq] = helper_atioq
    ERFA_HELPERS[erfa_ufunc.atoiq] = helper_atoiq
    return ERFA_HELPERS


UFUNC_HELPERS.register_module('erfa.ufunc', erfa_ufuncs,
                              get_erfa_helpers)
