"""Some join classes for determining if entries are present in the database or
have been updated.
"""
from dmd import orm
from dmd.warehouse import orm as worm
from multi_join import join
import heapq


_VMP_NAME_FIELDS = ['vmp_name', 'vmp_name_abbrev', 'vmp_name_prev']
_AMP_NAME_FIELDS = ['amp_name', 'amp_desc', 'amp_name_abbrev', 'amp_name_prev']


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_amp_ids(session):
    """Yield the current VMP IDs from the DM+D

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session for interacting with the DM+D database.

    Yields
    ------
    amp_id : `int`
        The virtual medicinal product ID (could be current or previous)
    vmp_data : `dmd.orm.ActMedProd`
        The actual medicinal product data.
    """
    q = session.query(orm.ActMedProd).order_by(orm.ActMedProd.amp_id)
    for i in q:
        yield i.amp_id, i


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_current_vmp_ids(session):
    """Yield the current VMP IDs from the DM+D

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session for interacting with the DM+D database.

    Yields
    ------
    vmp_id : `int`
        The virtual medicinal product ID (could be current or previous)
    vmp_data : `dmd.orm.VirtMedProd`
        The virtual virtual product data.
    """
    q = session.query(orm.VirtMedProd).order_by(orm.VirtMedProd.vmp_id)
    for i in q:
        yield i.vmp_id, i


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_dmd_vmp_names(session):
    """Yield the vmp names from the DM+D, grouped with a hash of all the names

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session for interacting with the DM+D database.

    Yields
    ------
    vmp_id : `int`
        The virtual medicinal product ID (could be current or previous)
    vmp_name_hash : `int`
        A python hash of the sorted VMP names.
    vmps : `list` of `dmd.orm.VirtMedProd`
        The virtual medicinal product data represented in the name hash.
    """
    q = session.query(orm.VirtMedProd).order_by(orm.VirtMedProd.vmp_id)
    for i in _hash_query(iter(q), lambda x: x.vmp_id, _extract_vmp_names):
        yield i


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_dmd_amp_names(session):
    """Yield the vmp names from the DM+D, grouped with a hash of all the names

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session for interacting with the DM+D database.

    Yields
    ------
    vmp_id : `int`
        The virtual medicinal product ID (could be current or previous)
    vmp_name_hash : `int`
        A python hash of the sorted VMP names.
    vmps : `list` of `dmd.orm.VirtMedProd`
        The virtual medicinal product data represented in the name hash.
    """
    q = session.query(orm.ActMedProd).order_by(orm.ActMedProd.amp_id)
    for i in _hash_query(iter(q), lambda x: x.amp_id, _extract_amp_names):
        yield i


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_wh_dmd_names(session):
    """Yield the vmp names from the DM+D, grouped with a hash of all the names

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session for interacting with the DM+D database.

    Yields
    ------
    vmp_id : `int`
        The virtual medicinal product ID (could be current or previous)
    vmp_name_hash : `int`
        A python hash of the sorted VMP names.
    vmps : `list` of `dmd.orm.VirtMedProd`
        The virtual medicinal product data represented in the name hash.
    """
    q = session.query(worm.DmdName).\
        filter(worm.DmdName.is_latest == True).\
        order_by(worm.DmdName.dmd_id)

    for i in _hash_query(iter(q), lambda x: x.dmd_id, lambda x: [x.dmd_name]):
        yield i



# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_dmd_vmp_unit_dose(session):
    """Yield the vmp names from the DM+D, grouped with a hash of all the names

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session for interacting with the DM+D database.

    Yields
    ------
    vmp_id : `int`
        The virtual medicinal product ID (could be current or previous)
    vmp_name_hash : `int`
        A python hash of the sorted VMP names.
    vmps : `list` of `dmd.orm.VirtMedProd`
        The virtual medicinal product data represented in the name hash.
    """
    q = session.query(orm.VirtMedProd).order_by(orm.VirtMedProd.vmp_id)
    for i in _hash_query(iter(q), lambda x: x.vmp_id, _extract_vmp_unit_dose):
        yield i


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_dmd_amp_unit_dose(session):
    """Yield the vmp names from the DM+D, grouped with a hash of all the names

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session for interacting with the DM+D database.

    Yields
    ------
    vmp_id : `int`
        The virtual medicinal product ID (could be current or previous)
    vmp_name_hash : `int`
        A python hash of the sorted VMP names.
    vmps : `list` of `dmd.orm.VirtMedProd`
        The virtual medicinal product data represented in the name hash.
    """
    q = session.query(orm.ActMedProd).order_by(orm.ActMedProd.amp_id)
    for i in _hash_query(iter(q), lambda x: x.amp_id, _extract_amp_unit_dose):
        yield i


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_wh_dmd_unit_dose(session):
    """Yield the vmp names from the DM+D, grouped with a hash of all the names

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session for interacting with the DM+D database.

    Yields
    ------
    vmp_id : `int`
        The virtual medicinal product ID (could be current or previous)
    vmp_name_hash : `int`
        A python hash of the sorted VMP names.
    vmps : `list` of `dmd.orm.VirtMedProd`
        The virtual medicinal product data represented in the name hash.
    """
    q = session.query(worm.UnitDose).\
        filter(worm.UnitDose.is_latest == True).\
        order_by(worm.UnitDose.dmd_id)

    for i in _hash_query(
            iter(q), lambda x: x.dmd_id,
            lambda x: [
                (
                    x.dose_form_ind,
                    x.unit_dose_form_size,
                    x.unit_dose_form_units,
                    x.unit_dose_form_type
                 )
            ]
    ):
        yield i


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_dmd_vmp_ingredients(session):
    """Yield the vmp names from the DM+D, grouped with a hash of all the names

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session for interacting with the DM+D database.

    Yields
    ------
    vmp_id : `int`
        The virtual medicinal product ID (could be current or previous)
    vmp_name_hash : `int`
        A python hash of the sorted VMP names.
    vmps : `list` of `dmd.orm.VirtMedProd`
        The virtual medicinal product data represented in the name hash.
    """
    q = session.query(orm.VirtMedProd).order_by(orm.VirtMedProd.vmp_id)
    for i in _hash_query(iter(q), lambda x: x.vmp_id, _extract_vmp_ingredients):
        yield i


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_dmd_amp_ingredients(session):
    """Yield the vmp names from the DM+D, grouped with a hash of all the names

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session for interacting with the DM+D database.

    Yields
    ------
    vmp_id : `int`
        The virtual medicinal product ID (could be current or previous)
    vmp_name_hash : `int`
        A python hash of the sorted VMP names.
    vmps : `list` of `dmd.orm.VirtMedProd`
        The virtual medicinal product data represented in the name hash.
    """
    q = session.query(orm.ActMedProd).order_by(orm.ActMedProd.amp_id)
    for i in _hash_query(iter(q), lambda x: x.amp_id, _extract_amp_ingredients):
        yield i


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_wh_dmd_ingredients(session):
    """Yield the vmp names from the DM+D, grouped with a hash of all the names

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session for interacting with the DM+D database.

    Yields
    ------
    vmp_id : `int`
        The virtual medicinal product ID (could be current or previous)
    vmp_name_hash : `int`
        A python hash of the sorted VMP names.
    vmps : `list` of `dmd.orm.VirtMedProd`
        The virtual medicinal product data represented in the name hash.
    """
    q = session.query(worm.Ingredient).\
        filter(worm.Ingredient.is_latest == True).\
        order_by(worm.Ingredient.dmd_id)

    for i in _hash_query(
            iter(q), lambda x: x.dmd_id,
            lambda x: [
                (
                    x.ing_id,
                    x.ing_name,
                    x.stren_val_num,
                    x.stren_val_num_unit,
                    x.stren_val_den,
                    x.stren_val_den_unit
                 )
            ]
    ):
        yield i


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _hash_query(query, group_func, hash_data_func):
    """
    """
    first_row = next(query)
    cur_id = group_func(first_row)
    hash_data = set(hash_data_func(first_row))
    rows = [first_row]

    for i in query:
        group_var = group_func(i)
        if group_var != cur_id:
            yield cur_id, hash(tuple(sorted(hash_data))), rows
            cur_id = group_var
            rows = []
            hash_data = set([])
        hash_data.update(hash_data_func(i))
        rows.append(i)

    if len(hash_data) > 0:
        yield cur_id, hash(tuple(sorted(hash_data))), rows

# # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# def get_vmp_id(row):
#     """
#     """
#     return row.vmp_id


# # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# def get_dmd_name(row):
#     """
#     """
#     return [row.dmd_name]

# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _extract_vmp_names(row):
    """
    """
    names = []
    for i in _VMP_NAME_FIELDS:
        n = getattr(row, i)
        if n is not None:
            names.append(n)
    return names


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _extract_amp_names(row):
    """
    """
    names = []
    for i in _AMP_NAME_FIELDS:
        n = getattr(row, i)
        if n is not None:
            names.append(n)
    return names


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _extract_vmp_unit_dose(row):
    """Get ``UnitDose`` objects for each of the defined names in a AMP/VMP.

    Parameters
    ----------
    dmd_code : `dmd.warehouse.orm.DmdCode`
        The code object.
    vmp : `func`
        The virtual medical product relating to the dmd_code.
    version : `dmd.warehouse.orm.DmdVersion`, optional, default: `None`
        The version object, if not provided then it is not added.
    is_latest : `bool`, optional, default: `False`
        The ``is_latest`` flag.

    Returns
    -------
    unit_dose : ``dmd.warehouse.orm.UnitDose`
        The unit dose object.
    """
    try:
        unit_dose_form_type = row.unit_meas_ud.unit_meas_desc
    except AttributeError:
        unit_dose_form_type = None

    try:
        unit_dose_form_units = row.unit_meas.unit_meas_desc
    except AttributeError:
        unit_dose_form_units = None

    return [
        (
            row.dose_form_ind.dose_form_ind_desc,
            row.unit_dose_form_size,
            unit_dose_form_units,
            unit_dose_form_type
        )
    ]


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _extract_amp_unit_dose(row):
    return _extract_vmp_unit_dose(row.virt_med_prod)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _extract_vmp_ingredients(row):
    """Get ``UnitDose`` objects for each of the defined names in a AMP/VMP.

    Parameters
    ----------
    dmd_code : `dmd.warehouse.orm.DmdCode`
        The code object.
    vmp : `func`
        The virtual medical product relating to the dmd_code.
    version : `dmd.warehouse.orm.DmdVersion`, optional, default: `None`
        The version object, if not provided then it is not added.
    is_latest : `bool`, optional, default: `False`
        The ``is_latest`` flag.

    Returns
    -------
    unit_dose : ``dmd.warehouse.orm.UnitDose`
        The unit dose object.
    """
    return [tuple(i.values()) for i in get_vmp_ingredients(row)]


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _extract_amp_ingredients(row):
    return _extract_vmp_ingredients(row.virt_med_prod)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_vmp_ingredients(row):
    ingredients = []
    vpis = row.virt_prod_ing
    if not isinstance(vpis, list):
        print(vpis)
        print(type(vpis))
        raise TypeError("expecting list")

    for i in vpis:
        try:
            stren_val_num_unit = i.stren_value_num_unit_meas.unit_meas_desc
        except AttributeError:
            stren_val_num_unit = None

        try:
            stren_val_den_unit = i.stren_value_den_unit_meas.unit_meas_desc
        except AttributeError:
            stren_val_den_unit = None

        ingredients.append(
            dict(
                ing_id=i.vmp_ingredient_id,
                ing_name=i.ingredient.ingredient_name,
                stren_val_num=i.stren_value_num,
                stren_val_num_unit=stren_val_num_unit,
                stren_val_den=i.stren_value_den,
                stren_val_den_unit=stren_val_den_unit
            )
        )
    return ingredients


# # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
# def yield_prev_vmp_ids(session):
#     """Yield the VMP IDs (or previous IDs) from the DM+D
#     """
#     q = session.query(orm.VirtMedProd).\
#         filter(orm.VirtMedProd.vmp_id_prev != None).\
#         group_by(orm.VirtMedProd.vmp_id_prev).\
#         order_by(orm.VirtMedProd.vmp_id_prev)

#     for i in q:
#         yield [i.vmp_id_prev, i]


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_all_vmp_ids(session):
    """Yield all VMP IDs (or previous IDs) from the DM+D.

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session for interacting with the DM+D database.

    Yields
    ------
    vmp_id : `int`
        The virtual medicinal product ID (could be current or previous)
    vmp_data : `dmd.orm.VirtMedProd`
        The virtual virtual product data.

    Notes
    -----
    This essentially does a union of the current and previous VMP IDs. However,
    I could not get union working properly in SQLAlchemy so I am improvising by
    using ``heapq.merge``.
    """
    label = 'join_id'
    q1 = session.query(orm.VirtMedProd.vmp_id.label(label),
                       orm.VirtMedProd).\
        order_by(orm.VirtMedProd.vmp_id)

    q2 = session.query(orm.VirtMedProd.vmp_id_prev.label(label),
                       orm.VirtMedProd).\
        filter(orm.VirtMedProd.vmp_id_prev != None).\
        group_by(orm.VirtMedProd.vmp_id_prev).\
        order_by(orm.VirtMedProd.vmp_id_prev)

    # Get SQLite errors with this so using heapq.merge
    # q1.union(q2).order_by(label)
    # Works as both queries have an order by statement
    for i in heapq.merge(q1, q2, key=lambda x: x[0]):
        yield i


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_dmd_ids(session, is_amp=None):
    """Yield the DMD IDs (VMP or AMP IDs) from the warehouse (in code order).

    Parameters
    ----------
    session : `sqlalchemy.Session`
        A session for interacting with the DM+D warehouse database.
    id_amp : `bool` or `NoneType`, optional, default: `NoneType`
        Query those records that are VMPs (``False``), AMPs (``True``) or
        all (``NoneType``).

    Yields
    ------
    dmd_id : `int`
        The dmd ID
    dmd_code_data : `dmd.warehouse.orm.DmdCode`
        The DM+D code data.
    """
    q = session.query(worm.DmdCode)

    if is_amp is not None:
        q = q.filter(worm.DmdCode.is_amp == is_amp)

    q = q.order_by(worm.DmdCode.dmd_id)
    for i in q:
        yield i.dmd_id, i


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def join_on_ids(wh_query, dmd_query):
    """Perform a join between the warehouse database query and the DM+D
    database query.

    Parameters
    ----------
    wh_query : ``
        The query against the DM+D database. Must yield rows with the join ID
        at [0].
    dmd_query : ``
        The query against the DM+D database. Must yield rows with the join ID
        at [0].

    Yields
    ------
    joined_row : `join.`
        A joined row. Please note, that in the join both queries are set to
        return all their results so, some of the rows will be asymmetric,
        think of this like a full join.
    """
    dmd = join.JoinIterator(
        dmd_query, key=(0,),
        monitor=None, join_condition=join.RETURN_ALL
    )
    wh = join.JoinIterator(
        wh_query, key=(0,),
        monitor=None, join_condition=join.RETURN_ALL
    )
    j = join.Join(wh, dmd)
    for i in j:
        yield i
