import astropy as ap
from astropy.units import Quantity
from astropy.constants import Constant, EMConstant
from astropy.table import Table

from hickle.helpers import get_type_and_data
import six

def create_astropy_quantity(py_obj, h_group, call_id=0, **kwargs):
    """ dumps an astropy quantity

    Args:
        py_obj: python object to dump; should be a python type (int, float, bool etc)
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    # kwarg compression etc does not work on scalars
    d = h_group.create_dataset('data_%i' % call_id, data=py_obj.value,
                               dtype='float64')     #, **kwargs)
    d.attrs["type"] = [b'astropy_quantity']
    if six.PY3:
        unit = bytes(str(py_obj.unit), 'ascii')
    else:
        unit = str(py_obj.unit)
    d.attrs['unit'] = [unit]


def create_astropy_constant(py_obj, h_group, call_id=0, **kwargs):
    """ dumps an astropy constant

    Args:
        py_obj: python object to dump; should be a python type (int, float, bool etc)
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    # kwarg compression etc does not work on scalars
    d = h_group.create_dataset('data_%i' % call_id, data=py_obj.value,
                               dtype='float64')     #, **kwargs)
    d.attrs["type"]   = [b'astropy_constant']
    d.attrs["unit"]   = [str(py_obj.unit)]
    d.attrs["abbrev"] = [str(py_obj.abbrev)]
    d.attrs["name"]   = [str(py_obj.name)]
    d.attrs["reference"] = [str(py_obj.reference)]
    d.attrs["uncertainty"] = [py_obj.uncertainty]

    if py_obj.system:
        d.attrs["system"] = [py_obj.system]


def create_astropy_table(py_obj, h_group, call_id=0, **kwargs):
    """ Dump an astropy Table

    Args:
        py_obj: python object to dump; should be a python type (int, float, bool etc)
        h_group (h5.File.group): group to dump data into.
        call_id (int): index to identify object's relative location in the iterable.
    """
    data = py_obj.as_array()
    d = h_group.create_dataset('data_%i' % call_id, data=data, dtype=data.dtype, **kwargs)
    d.attrs['type']  = [b'astropy_table']

    if six.PY3:
        colnames = [bytes(cn, 'ascii') for cn in py_obj.colnames]
    else:
        colnames = py_obj.colnames
    d.attrs['colnames'] = colnames
    for key, value in py_obj.meta.items():
     d.attrs[key] = value


def load_astropy_quantity_dataset(h_node):
    py_type, data = get_type_and_data(h_node)
    unit = h_node.attrs["unit"][0]
    q = Quantity(data, unit)
    return q

def load_astropy_constant_dataset(h_node):
    py_type, data = get_type_and_data(h_node)
    unit   = h_node.attrs["unit"][0]
    abbrev = h_node.attrs["abbrev"][0]
    name   = h_node.attrs["name"][0]
    ref    = h_node.attrs["reference"][0]
    unc    = h_node.attrs["uncertainty"][0]

    system = None
    if "system" in h_node.attrs.keys():
        system = h_node.attrs["system"][0]

    c = Constant(abbrev, name, data, unit, unc, ref, system)
    return c

def load_astropy_table(h_node):
    py_type, data = get_type_and_data(h_node)
    metadata = dict(h_node.attrs.items())
    metadata.pop('type')
    metadata.pop('colnames')

    if six.PY3:
        colnames = [cn.decode('ascii') for cn in h_node.attrs["colnames"]]
    else:
        colnames = h_node.attrs["colnames"]

    t = Table(data, names=colnames, meta=metadata)
    return t

def check_is_astropy_table(py_obj):
    return isinstance(py_obj, Table)

#######################
## Lookup dictionary ##
#######################

types_dict = {
    Quantity: create_astropy_quantity,
    Constant: create_astropy_constant,
    Table: create_astropy_table
}

hkl_types_dict = {
    b'astropy_quantity' : load_astropy_quantity_dataset,
    b'astropy_constant' : load_astropy_constant_dataset,
    b'astropy_table'    : load_astropy_table
}