"""Some simple pre-formulated queries, to use against DM+D
"""
from sqlalchemy import or_
from dmd import (
    orm as o,
    common
)
from dmd.units import Unit
from tqdm import tqdm
import re


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class Ingredient(object):
    """A representation of a product ingredient.

    Parameters
    ----------
    name : `str`
        The ingredient name
    unit : `dmd.query.Unit`, optional, default: `None`
        The ingredient strength, and unit of measure.

    Notes
    -----
    This is used to represent the ``Product`` ``Ingredients`` in the cache.
    """
    STRENGTH_REGEXP = re.compile(r'^(?P<VALUE>\d+(?:\.\d*)?)')
    """A regular expression used to extract floats or ints from in front of
    units (`re.Pattern`)
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, name, unit=None):
        self.name = name.strip()
        self.query_name = self.name.lower()
        self.unit = unit

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __repr__(self):
        """Pretty printing
        """
        # Build in the attributes
        attrs = [
            "name='{0}'".format(self.name),
            "query_name='{0}'".format(self.query_name),
            "unit='{0}'".format(self.unit)
        ]
        return "<{0}({1})>".format(self.__class__.__name__, ",".join(attrs))

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    @classmethod
    def parse_strength(cls, strength):
        """Parse the strength into it's component parts, such as numerator
        value, numerator units, denominator value and denominator units.

        Parameters
        ----------
        strength : `str`
            The strength string to parse.

        Returns
        -------
        num_value : `float`
            The value of the numerator.
        num_units : `str` or `NoneType`
            The units of the numerator.
        den_value : `float` or `NoneType`
            The value of the denominator.
        den_units : `str` or `NoneType`
            The units of the denominator.
        """
        if strength is None or strength == '':
            return None, None, None, None

        # For things like: '50microgram/24hr (3.2mg/unit)'
        strength = re.sub(
            r'(.+?)\s*(?P<OTHER>\(.+?\))?$', '\\1',
            strength.strip()
        ).strip()

        # '2000iu (16.8 Microgram)/0.5ml'
        strength = re.sub(
            r'\d+\s*i?(?:u|units)\s*\(\s*(?P<FIRST>\d+(?:\.\d+)?\s*\w+)\)(?P<SECOND>.+)?$', '\\1\\2',
            strength
        )
        # '250mg + 250,000 Units'
        strength = re.sub(r',(?=\d)', '', strength)
        strength_match = cls.STRENGTH_REGEXP.match(strength)

        try:
            num_value = float(strength_match.group('VALUE'))
        except AttributeError:
            num_value = None

        units = \
            [i.strip()
             for i in cls.STRENGTH_REGEXP.sub('', strength).split('/')]
        try:
            num_units = units[0]
        except IndexError:
            num_units = None

        try:
            den_val_match = cls.STRENGTH_REGEXP.match(units[1])
            den_value = float(den_val_match.group('VALUE'))
            den_units = cls.STRENGTH_REGEXP.sub('', units[1])
        except IndexError:
            den_value = None
            den_units = None
        except AttributeError:
            den_value = 1
            den_units = units[1]

        return num_value, num_units, den_value, den_units


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class DmdEntry(object):
    """A single DM+D product entry (VMP or AMP).

    Wraps identity, naming, ingredient, and unit-dose information for
    a virtual or actual medicinal product.  Instances are hashable and
    compared by ``dmd_id``.

    Parameters
    ----------
    dmd_id : `int`
        The DM+D identifier for this product.
    dmd_name : `str`
        The product name.
    is_amp : `bool`
        ``True`` if this entry is an AMP, ``False`` for a VMP.
    vmp_id : `int`
        The parent VMP identifier (same as ``dmd_id`` for VMPs).
    db_version : `str`
        The DM+D database version this entry was loaded from.
    ingredients : `list` of :class:`Ingredient`, optional
        Active ingredients for this product.
    unit_dose_form_size : `str` or ``None``, optional
        Unit dose form size description, if available.
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, dmd_id, dmd_name, is_amp, vmp_id, db_version,
                 ingredients=None, unit_dose_form_size=None):
        self.dmd_id = dmd_id
        self.dmd_name = dmd_name
        self.is_amp = is_amp
        self.db_version = db_version
        self.vmp_id = vmp_id
        self.unit_dose_form_size = unit_dose_form_size
        self.ingredients = ingredients or list()

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __hash__(self):
        return hash(self.dmd_id)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __eq__(self, other):
        return self.dmd_id == other.dmd_id

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __ne__(self, other):
        return self.dmd_id != other.dmd_id

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __repr__(self):
        attrs = []
        for i in ['dmd_id', 'dmd_name', 'is_amp', 'vmp_id', 'db_version',
                  'ingredients']:
            attrs.append("{0}={1}".format(i, getattr(self, i)))
        return "<{0}({1})".format(self.__class__.__name__, ",".join(attrs))


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def create_caches(sms, verbose=True):
    """
    """
    id_cache = dict()
    name_cache = dict()
    ing_cache = dict()
    for i in tqdm(sms, disable=not verbose, unit=" databases",
                  desc='[info] building cache'):
        session = i()
        db_version = session.query(o.FileInfo.file_basename).filter(
            o.FileInfo.file_type == 'zip'
        ).one()
        db_version = re.sub(r'\.zip$', '', db_version[0])
        for row in session.query(o.ActMedProd):
            entries = _create_amp_entries(row, db_version)
            _update_id_cache(entries, id_cache)
            _update_name_cache(entries, name_cache)
            _update_ing_cache(entries, ing_cache)

        for row in session.query(o.VirtMedProd):
            entries = _create_vmp_entries(row, db_version)
            _update_id_cache(entries, id_cache)
            _update_name_cache(entries, name_cache)
            _update_ing_cache(entries, ing_cache)

        session.close()
    return id_cache, name_cache, ing_cache


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _create_amp_entries(row, version):
    """
    """
    vmp = row.virt_med_prod
    unit_dose_form_size = vmp.unit_dose_form_size
    ings = vmp.virt_prod_ing
    all_ings = _parse_ings(ings)
    dmd_id = row.amp_id
    vmp_id = row.vmp_id

    all_entries = []
    for i in [row.amp_name, row.amp_desc, row.amp_name_abbrev,
              row.amp_name_prev]:
        if i is not None:
            all_entries.append(
                DmdEntry(dmd_id, i, True, vmp_id, version,
                         ingredients=all_ings,
                         unit_dose_form_size=unit_dose_form_size)
            )
    return all_entries


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _create_vmp_entries(row, version):
    """
    """
    ings = row.virt_prod_ing
    all_ings = _parse_ings(ings)
    vmp_id = row.vmp_id
    unit_dose_form_size = row.unit_dose_form_size
    all_entries = []

    for i in [row.vmp_id, row.vmp_id_prev]:
        if i is not None:
            dmd_id = i
            for j in [row.vmp_name, row.vmp_name_abbrev, row.vmp_name_prev]:
                if j is not None:
                    all_entries.append(
                        DmdEntry(dmd_id, j, False, vmp_id, version,
                                 ingredients=all_ings,
                                 unit_dose_form_size=unit_dose_form_size)
                    )
    return all_entries


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _parse_ings(ings):
    """
    """
    all_ings = []
    for i in ings:
        try:
            num_unit = i.stren_value_num_unit_meas.unit_meas_desc
        except AttributeError:
            num_unit = None

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

        unit = Unit(
            i.stren_value_num,
            num_unit=num_unit,
            den_value=i.stren_value_den,
            den_unit=den_unit
        )
        all_ings.append(
            Ingredient(
                i.ingredient.ingredient_name, unit=unit
            )
        )
    return all_ings


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _update_id_cache(entries, id_cache):
    """Add entries to the ID cache, skipping duplicates.

    Parameters
    ----------
    entries : iterable of :class:`DmdEntry`
        Entries to index.
    id_cache : `dict`
        Mapping of ``dmd_id`` -> :class:`DmdEntry`.
    """
    for entry in entries:
        try:
            id_cache[entry.dmd_id].dmd_id
        except KeyError:
            id_cache[entry.dmd_id] = entry


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _update_name_cache(entries, name_cache):
    """Add entries to the name cache, grouping by lowercase name.

    Parameters
    ----------
    entries : iterable of :class:`DmdEntry`
        Entries to index.
    name_cache : `dict`
        Mapping of lowercase name -> list of :class:`DmdEntry`.
    """
    for entry in entries:
        try:
            cur_ids = [n.dmd_id for n in name_cache[entry.dmd_name.lower()]]
            if entry.dmd_id not in cur_ids:
                name_cache[entry.dmd_name.lower()].append(entry)
        except KeyError:
            name_cache[entry.dmd_name.lower()] = [entry]


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _update_ing_cache(entries, ing_cache):
    """Add entries to the ingredient cache, keyed by ingredient name.

    Parameters
    ----------
    entries : iterable of :class:`DmdEntry`
        Entries to index (must have an ``ingredients`` attribute).
    ing_cache : `dict`
        Mapping of ingredient query name -> set of :class:`DmdEntry`.
    """
    for entry in entries:
        for ing in entry.ingredients:
            try:
                ing_cache[ing.query_name].add(entry)
            except KeyError:
                ing_cache[ing.query_name] = set([entry])


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_amp_by_id(session, amp_id):
    """Get the actual medicinal products from the AMP ID

    Parameters
    ----------
    session : `sqlalchemy.Session`
        For interaction with the database
    amp_id  : `int`
        The actual medicinal product identifier

    Returns
    -------
    matches : `list` of `dmd.orm.ActMedProd`
        This list will have length 0 if there are no matches.
    """
    return session.query(o.ActMedProd).filter(
        o.ActMedProd.amp_id == amp_id
    ).all()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_vmp_by_id(session, vmp_id):
    """Get the virtual medicinal products from the VMP ID

    Parameters
    ----------
    session : `sqlalchemy.Session`
        For interaction with the database
    vmp_id  : `int`
        The virtual medicinal product identifier

    Returns
    -------
    matches : `list` of `dmd.orm.VirtMedProd`
        This list will have length 0 if there are no matches.
    """
    return session.query(o.VirtMedProd).filter(
        o.VirtMedProd.vmp_id == vmp_id
    ).all()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_amp_by_name(session, amp_name):
    """Get the actual medicinal products from the AMP name

    Parameters
    ----------
    session : `sqlalchemy.Session`
        For interaction with the database
    amp_id  : `int`
        The actual medicinal product identifier

    Returns
    -------
    matches : `list` of `dmd.orm.ActMedProd`
        This list will have length 0 if there are no matches.

    Notes
    -----
    This will query for ``amp_name``, ``amp_name_abbrev``, ``amp_desc``,
    ``amp_name_prev``.
    """
    return session.query(o.ActMedProd).filter(
        or_(
            o.ActMedProd.amp_name == amp_name,
            o.ActMedProd.amp_name_abbrev == amp_name,
            o.ActMedProd.amp_desc == amp_name,
            o.ActMedProd.amp_name_prev == amp_name,
        )
    ).all()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_vmp_by_name(session, vmp_name):
    """Get the actual medicinal products from the VMP name

    Parameters
    ----------
    session : `sqlalchemy.Session`
        For interaction with the database
    vmp_name  : `int`
        The virtual medicinal product identifier

    Returns
    -------
    matches : `list` of `dmd.orm.VirtMedProd`
        This list will have length 0 if there are no matches.

    Notes
    -----
    This will query for ``vmp_name``, ``vmp_name_abbrev``,
    ``vmp_name_prev``.
    """
    return session.query(o.VirtMedProd).filter(
        or_(
            o.VirtMedProd.vmp_name == vmp_name,
            o.VirtMedProd.vmp_name_abbrev == vmp_name,
            o.VirtMedProd.vmp_name_prev == vmp_name,
        )
    ).all()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_ingredient_by_name(session, ingredient_name):
    """Get the actual medicinal products from the VMP name

    Parameters
    ----------
    session : `sqlalchemy.Session`
        For interaction with the database
    vmp_name  : `int`
        The virtual medicinal product identifier

    Returns
    -------
    matches : `list` of `dmd.orm.VirtMedProd`
        This list will have length 0 if there are no matches.

    Notes
    -----
    This will query for ``vmp_name``, ``vmp_name_abbrev``,
    ``vmp_name_prev``.
    """
    return session.query(o.Ingredient).filter(
        o.Ingredient.ingredient_name.like(ingredient_name)
    ).all()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_vtm_by_name(session, name):
    """Get the actual medicinal products from the VMP name

    Parameters
    ----------
    session : `sqlalchemy.Session`
        For interaction with the database
    vmp_name  : `int`
        The virtual medicinal product identifier

    Returns
    -------
    matches : `list` of `dmd.orm.VirtMedProd`
        This list will have length 0 if there are no matches.

    Notes
    -----
    This will query for ``vmp_name``, ``vmp_name_abbrev``,
    ``vmp_name_prev``.
    """
    return session.query(o.VirtTherMoiet).filter(
        o.VirtTherMoiet.vtm_name.like(name)
    ).all()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_version_day(session):
    """Return the version day number. This is the number of days since the
    1/1/2000

    Parameters
    ----------
    session : `sqlalchemy.Session`
        For interaction with the database

    Returns
    -------
    db_days : `int`
        The number of days post 1/1/2000 that the DM+D database was released.
    """
    db_date = session.query(o.FileInfo.file_date).\
        filter(o.FileInfo.file_type == 'zip').one()
    db_days = (db_date[0] - common.DMD_EPOCH_DATE).days
    return db_days


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_version_name(session):
    """Return the version day number. This is the number of days since the
    1/1/2000

    Parameters
    ----------
    session : `sqlalchemy.Session`
        For interaction with the database

    Returns
    -------
    db_days : `int`
        The number of days post 1/1/2000 that the DM+D database was released.
    """
    db_name = session.query(o.FileInfo.file_basename).\
        filter(o.FileInfo.file_type == 'zip').one()
    return re.sub(r'\.zip$', '', db_name[0])
