import warnings
import re

from lxml import etree
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship, backref, object_session
from sqlalchemy.ext.associationproxy import association_proxy
from sqlalchemy import (Numeric, Unicode,
                        Column, Integer, ForeignKey,
                        UnicodeText, Boolean, event)
from sqlalchemy import exc as sa_exc
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from glypy.composition.composition import Composition, _make_isotope_string

Base = declarative_base()

_unimod_xml_download_url = "http://www.unimod.org/xml/unimod_tables.xml"

try:
    basestring
except NameError:
    basestring = (str, bytes)


CompositionType = Composition


def simple_repr(self):  # pragma: no cover
    template = "{self.__class__.__name__}({d})"
    items = [(k, v) for k, v in self.__dict__.items() if not k.startswith("_")]
    parts = []
    for key, value in items:
        if key.endswith("_id"):
            try:
                value = getattr(self, key[:-3])
                key = key[:-3]
            except AttributeError:
                pass
        parts.append("%s=%r" % (key, value))

    return template.format(self=self, d=', '.join(parts))

Base.__repr__ = simple_repr


def attr_repr(attr_name):  # pragma: no cover
    def __repr__(self):
        return str(getattr(self, attr_name))

    def applier(cls):
        cls.__repr__ = __repr__
        return cls

    return applier


def remove_namespace(doc, namespace):
    """Remove namespace in the passed document in place."""
    ns = u'{%s}' % namespace
    nsl = len(ns)
    for elem in doc.getiterator():
        if elem.tag.startswith(ns):
            elem.tag = elem.tag[nsl:]


def preprocess_xml(doc_path):
    '''
    Parse and drop namespaces from an XML document.

    Parameters
    ----------
    doc_path: str

    Returns
    -------
    etree.ElementTree
    '''
    tree = etree.parse(doc_path)
    root = tree.getroot()
    for ns in root.nsmap.values():
        remove_namespace(tree, ns)
    return tree


def _formula_parser(formula, session):
    '''
    Parse a unimod formula composed of elements,
    isotopes, and other bricks.

    In order to look up a Brick's composition, this
    function must have access to a session.

    Parameters
    ----------
    formula: str
        A Unimod formula of the form `A(n) B(m)...`
        where A, B, ... are element names or bricks and
        (n), (m)... are parenthesized possibly signed integers or
        omitted in which case they are interpreted as 1
    session: Session
        An active SQLAlchemy session for looking up bricks in the database

    Returns
    -------
    CompositionType
    '''
    composition = CompositionType()
    for token in formula.split(" "):
        match = re.search(r"(?P<isotope>\d+)?(?P<elemet>[^\(]+)(?:\((?P<count>-?\d+)\))?", token)
        if match:
            isotope, element, count = match.groups()
            if count is not None:
                count = int(count)
            else:
                count = 1
            if isotope is not None:
                name = _make_isotope_string(element, int(isotope))
            else:
                name = element
            if name == '0':
                continue
            is_brick = session.query(Brick).filter(Brick.brick == name).first()
            if is_brick is None:
                composition[str(name)] += count
            else:
                composition += is_brick.composition * count
    return composition


def _composition_listener(attr):
    '''
    Attach event listeners to an InstrumentedAttribute
    to trigger formula parsing on load and on change.
    '''
    @event.listens_for(attr, "set")
    def _update_composition_from_formula(target, value, oldvalue, initiator):
        session = object_session(target)
        if value == "" or value is None:
            return
        # If the object hasn't been associated with a session,
        # we can't look up bricks.
        if session is None:
            return
        target.composition = _formula_parser(value, session)

    @event.listens_for(attr.class_, "load")
    def _update_composition_on_load(target, context):
        value = getattr(target, attr.prop.key)
        if value == "" or value is None:
            return
        session = object_session(target)
        target.composition = _formula_parser(value, session)


def has_composition(attr_name):
    '''
    A decorator to simplify flagging a Model with a column
    to be treated as a formula for parsing. Calls :func:`_composition_listener`
    internally.
    '''
    def decorator(model):
        _composition_listener(getattr(model, attr_name))
        return model
    return decorator


class HasFullNameMixin(object):
    '''
    A simple mixin to standardize equality operators
    for models with a :attr:`full_name` attribute.
    '''
    def __eq__(self, other):
        try:
            return self.full_name == other.full_name
        except AttributeError:
            return self.full_name == other

    def __ne__(self, other):
        return not self == other

    def __hash__(self):
        return hash(self.full_name)


class AlternativeName(Base):
    __tablename__ = "AlternativeName"

    _tag_name = "alt_names_row"

    @classmethod
    def from_tag(cls, tag):
        attrib = tag.attrib
        inst = cls(
            id=int(attrib["record_id"]),
            alt_name=attrib['alt_name'],
            modification_id=int(attrib['mod_key'])
            )
        return inst

    id = Column(Integer, primary_key=True)
    alt_name = Column(Unicode(256), index=True)
    modification_id = Column(Integer, ForeignKey("Modification.id"), index=True)


class AminoAcid(Base, HasFullNameMixin):
    __tablename__ = "AminoAcid"

    _tag_name = "amino_acids_row"

    @classmethod
    def from_tag(cls, tag):
        attrib = tag.attrib
        inst = cls(
            id=int(attrib['record_id']),
            full_name=attrib['full_name'],
            one_letter=attrib['one_letter'],
            three_letter=attrib['three_letter'],
            num_H=int(attrib["num_H"]),
            num_O=int(attrib["num_O"]),
            num_C=int(attrib["num_C"]),
            num_N=int(attrib["num_N"]),
            num_S=int(attrib["num_S"]),
            )
        return inst

    id = Column(Integer, primary_key=True)
    num_H = Column(Integer)
    num_O = Column(Integer)
    num_C = Column(Integer)
    num_N = Column(Integer)
    num_S = Column(Integer)
    full_name = Column(Unicode(25), index=True)
    one_letter = Column(Unicode(10), index=True)
    three_letter = Column(Unicode(10), index=True)


@attr_repr("classification")
class Classification(Base):
    __tablename__ = "Classification"

    _tag_name = "classifications_row"

    @classmethod
    def from_tag(cls, tag):
        attrib = tag.attrib
        inst = cls(
            id=int(attrib['record_id']),
            classification=attrib["classification"]
            )
        return inst

    id = Column(Integer, primary_key=True)
    classification = Column(Unicode(30), index=True)


@attr_repr("position")
class Position(Base):
    __tablename__ = "Position"

    _tag_name = "positions_row"

    @classmethod
    def from_tag(cls, tag):
        attrib = tag.attrib
        inst = cls(
            id=int(attrib['record_id']),
            position=attrib['position']
            )
        return inst

    def __eq__(self, other):
        try:
            return self.position == other.position
        except AttributeError:
            return self.position == other

    def __ne__(self, other):
        return not self == other

    def __hash__(self):
        return hash(self.position)

    id = Column(Integer, primary_key=True)
    position = Column(Unicode(20), index=True)


class Brick(Base, HasFullNameMixin):
    __tablename__ = "Brick"

    _tag_name = "bricks_row"

    @classmethod
    def from_tag(cls, tag):
        attrib = tag.attrib
        inst = cls(
            id=int(attrib['record_id']),
            brick=attrib['brick'],
            full_name=attrib['full_name']
            )
        return inst

    id = Column(Integer, primary_key=True)
    brick = Column(Unicode(64), index=True)
    full_name = Column(Unicode(128), index=True)

    elements = relationship("BrickToElement")

    @property
    def composition(self):
        composition = CompositionType()
        for element_relation in self.elements:
            symbol = element_relation.element
            isotope, element = re.search(r"(?P<isotope>\d*)?(?P<element>\S+)", symbol).groups()
            if isotope != "":
                isotope = int(isotope)
                iso_str = _make_isotope_string(element, isotope)
            else:
                iso_str = element
            count = element_relation.count
            composition[str(iso_str)] = count
        return composition


class Fragment(Base):
    __tablename__ = "Fragment"

    _tag_name = "fragments_row"

    @classmethod
    def from_tag(cls, tag):
        attrib = tag.attrib
        inst = cls(
            id=int(attrib['record_id']),
            modification_id=int(attrib["mod_key"])
            )
        return inst

    id = Column(Integer, primary_key=True)
    modification_id = Column(Integer, ForeignKey("Modification.id"), index=True)

    _fragment_composition = relationship("FragmentComposition")

    @property
    def composition(self):
        composition = CompositionType()
        session = object_session(self)
        for fragment_composition_relation in self._fragment_composition:
            symbol = fragment_composition_relation.brick_string
            isotope, element = re.search(r"(?P<isotope>\d*)?(?P<element>\S+)", symbol).groups()
            count = fragment_composition_relation.count
            if count is not None:
                count = int(count)
            else:
                count = 1
            if isotope != "":
                name = _make_isotope_string(element, isotope)
            else:
                name = element
            is_brick = session.query(Brick).filter(Brick.brick == name).first()
            if is_brick is None:
                composition[str(name)] += count
            else:
                composition += is_brick.composition * count
        return composition


class FragmentComposition(Base):
    __tablename__ = "FragmentComposition"

    _tag_name = "fragment_comp_row"

    @classmethod
    def from_tag(cls, tag):
        attrib = tag.attrib
        inst = cls(
            id=int(attrib['record_id']),
            brick_string=attrib["brick"],
            fragment_id=int(attrib["fragments_key"]),
            count=int(attrib["num_brick"])
            )
        return inst

    id = Column(Integer, primary_key=True)
    brick_string = Column(Unicode(64), ForeignKey(Brick.brick), index=True)
    fragment_id = Column(Integer, ForeignKey("Fragment.id"), index=True)
    count = Column(Integer)


class ModificationToBrick(Base):
    __tablename__ = "ModificationToBrick"

    _tag_name = "mod2brick_row"

    @classmethod
    def from_tag(cls, tag):
        attrib = tag.attrib
        inst = cls(
            id=int(attrib['record_id']),
            brick_string=(attrib['brick']),
            modification_id=int(attrib["mod_key"]),
            count=int(attrib["num_brick"])
            )
        return inst

    id = Column(Integer, primary_key=True)
    brick_string = Column(Unicode(64), ForeignKey(Brick.brick), index=True)
    modification_id = Column(Integer, ForeignKey("Modification.id"), index=True)
    count = Column(Integer)


class BrickToElement(Base):
    __tablename__ = "BrickToElement"

    _tag_name = "brick2element_row"

    @classmethod
    def from_tag(cls, tag):
        attrib = tag.attrib
        inst = cls(
            id=int(attrib['record_id']),
            brick_id=int(attrib['brick_key']),
            count=int(attrib["num_element"]),
            element=attrib['element']
            )
        return inst

    id = Column(Integer, primary_key=True)
    brick_id = Column(Integer, ForeignKey(Brick.id), index=True)
    element = Column(Unicode(16), ForeignKey("Element.element"), index=True)
    element_obj = relationship("Element", uselist=False)
    count = Column(Integer)


class Element(Base, HasFullNameMixin):
    __tablename__ = "Element"

    _tag_name = "elements_row"

    @classmethod
    def from_tag(cls, tag):
        attrib = tag.attrib
        inst = cls(
            id=int(attrib['record_id']),
            average_mass=float(attrib["avge_mass"]),
            monoisotopic_mass=float(attrib["mono_mass"]),
            full_name=attrib["full_name"],
            element=attrib["element"]

            )
        return inst

    id = Column(Integer, primary_key=True)
    average_mass = Column(Numeric(12, 6, asdecimal=False))
    monoisotopic_mass = Column(Numeric(12, 6, asdecimal=False))
    full_name = Column(Unicode(64), index=True)
    element = Column(Unicode(16), index=True)


@has_composition("_composition")
class Modification(Base, HasFullNameMixin):
    __tablename__ = "Modification"

    _tag_name = "modifications_row"

    id = Column(Integer, primary_key=True)
    username_of_poster = Column(Unicode(128))
    average_mass = Column(Numeric(12, 6, asdecimal=False), index=True)
    ex_code_name = Column(Unicode(64), index=True)
    monoisotopic_mass = Column(Numeric(12, 6, asdecimal=False), index=True)
    full_name = Column(Unicode(128), index=True)
    code_name = Column(Unicode(128), index=True)
    _composition = Column(Unicode(128), index=True)
    approved = Column(Boolean, index=True)

    notes = relationship("MiscNotesModifications")
    specificities = relationship("Specificity")
    bricks = relationship(ModificationToBrick)
    _fragments = relationship(Fragment)

    _alt_names = relationship(AlternativeName, backref=backref("modification"))
    # Maps the list of AlternativeName instances loaded dynamically from _alt_names
    # into a list of plain strings, since the AlternativeName type contains no
    # additional information.
    alternative_names = association_proxy("_alt_names", "alt_name")
    fragments = association_proxy("_fragments", "composition")

    @classmethod
    def from_tag(cls, tag):
        attrib = tag.attrib
        inst = cls(
            id=int(attrib["record_id"]),
            username_of_poster=attrib["username_of_poster"],
            average_mass=float(attrib["avge_mass"]),
            monoisotopic_mass=float(attrib["mono_mass"]),
            ex_code_name=attrib["ex_code_name"],
            code_name=attrib["code_name"],
            full_name=attrib["full_name"],
            approved=bool(int(attrib['approved'])),
            _composition=attrib["composition"]
            )
        for note in tag:
            if note.tag == MiscNotesModifications._tag_name:
                model_note = MiscNotesModifications._from_tag(note, inst.id)
                if model_note is not None:
                    inst.notes.append(model_note)
        return inst


class NoteBase(object):
    _tag_name = "misc_notes"

    id = Column(Integer, primary_key=True)
    text = Column(UnicodeText)

    def __str__(self):
        return self.text


class MiscNotesModifications(Base, NoteBase):
    __tablename__ = "MiscNotesModifications"

    modification_id = Column(Integer, ForeignKey(Modification.id), index=True)

    @classmethod
    def _from_tag(cls, tag, modification_id):
        if tag.text is None:
            return
        return cls(text=tag.text, modification_id=modification_id)


class Specificity(Base):
    __tablename__ = "Specificity"

    _tag_name = "specificity_row"

    id = Column(Integer, primary_key=True)
    position_id = Column(Integer, ForeignKey(Position.id), index=True)
    classification_id = Column(Integer, ForeignKey(Classification.id), index=True)
    classification = relationship("Classification", uselist=False)
    # Map through one_letter
    amino_acid = Column(Unicode(10), ForeignKey(AminoAcid.one_letter), index=True)
    modification_id = Column(Integer, ForeignKey(Modification.id), index=True)
    hidden = Column(Boolean, index=True)
    group = Column(Integer, index=True)
    neutral_losses = relationship("SpecificityToNeutralLoss")

    notes = relationship("MiscNotesSpecificities")

    position = relationship(Position)

    @classmethod
    def from_tag(cls, tag):
        attrib = tag.attrib
        inst = cls(
            id=int(attrib['record_id']),
            position_id=int(attrib['position_key']),
            classification_id=int(attrib["classifications_key"]),
            group=int(attrib['spec_group']),
            hidden=bool(int(attrib["hidden"])),
            amino_acid=attrib["one_letter"],
            modification_id=int(attrib["mod_key"]),
            )
        for note in tag:
            if note.tag == MiscNotesSpecificities._tag_name:
                model_note = MiscNotesSpecificities._from_tag(note, inst.id)
                if model_note is not None:
                    inst.notes.append(model_note)
        return inst


class MiscNotesSpecificities(Base, NoteBase):
    __tablename__ = "MiscNotesSpecificities"

    specificity_id = Column(Integer, ForeignKey(Specificity.id), index=True)

    @classmethod
    def _from_tag(cls, tag, specificity_id):
        if tag.text is None:
            return
        return cls(text=tag.text, specificity_id=specificity_id)


class NeutralLoss(Base):
    __tablename__ = "NeutralLoss"

    _tag_name = "neutral_losses_row"

    @classmethod
    def from_tag(cls, tag):
        attrib = tag.attrib
        inst = cls(
            id=int(attrib['record_id']),
            brick_string=(attrib['brick']),
            count=int(attrib["num_brick"]),
            specificity_id=int(attrib["spec_key"])
            )
        return inst

    id = Column(Integer, primary_key=True)
    brick_string = Column(Unicode(64), index=True)
    specificity_id = Column(Integer, ForeignKey(Specificity.id), index=True)
    count = Column(Integer)


@has_composition("_composition")
class SpecificityToNeutralLoss(Base):
    __tablename__ = "SpecificityToNeutralLoss"

    _tag_name = "spec2nl_row"

    @classmethod
    def from_tag(cls, tag):
        attrib = tag.attrib
        inst = cls(
            id=int(attrib['record_id']),
            specificity_id=int(attrib["spec_key"]),
            monoisotopic_mass=float(attrib["nl_mono_mass"]),
            average_mass=float(attrib["nl_avge_mass"]),
            is_required_peptide_neutral_loss=bool(int(attrib["is_req_pep_nl"])),
            is_peptide_neutral_loss=bool(int(attrib["is_pep_nl"])),
            is_slave=bool(int(attrib["is_slave_nl"])),
            _composition=attrib['nl_composition']
            )
        return inst

    id = Column(Integer, primary_key=True)
    specificity_id = Column(Integer, ForeignKey(Specificity.id), index=True)
    specificity = relationship(Specificity, uselist=False)
    monoisotopic_mass = Column(Numeric(12, 6, asdecimal=False), index=True)
    average_mass = Column(Numeric(12, 6, asdecimal=False), index=True)
    _composition = Column(Unicode(128))
    is_slave = Column(Boolean, index=True)
    is_peptide_neutral_loss = Column(Boolean, index=True)
    is_required_peptide_neutral_loss = Column(Boolean, index=True)


class CrossreferenceSource(Base):
    __tablename__ = "CrossreferenceSource"
    _tag_name = "xref_sources_row"

    id = Column(Integer, primary_key=True)
    source = Column(Unicode(64), index=True)

    @classmethod
    def from_tag(cls, tag):
        attrib = tag.attrib
        inst = cls()
        inst.id = int(attrib['record_id'])
        inst.source = attrib['xref_source']
        return inst


class Crossreference(Base):
    __tablename__ = "Crossreference"

    _tag_name = "xrefs_row"

    id = Column(Integer, primary_key=True)
    source_id = Column(Integer, ForeignKey(CrossreferenceSource.id), index=True)
    source = relationship(CrossreferenceSource, uselist=False)
    url = Column(Unicode(128))
    modification_id = Column(Integer, ForeignKey(Modification.id), index=True)
    text = Column(UnicodeText)

    @classmethod
    def from_tag(cls, tag):
        attrib = tag.attrib
        inst = cls()
        inst.id = int(attrib['record_id'])
        inst.url = attrib['xref_url']
        inst.source_id = int(attrib['xref_source_key'])
        inst.modification_id = int(attrib["mod_key"])
        text = []
        for node in tag.getchildren():
            if node.tag == "xref_text":
                if node.text is not None:
                    text.append(node.text)
        inst.text = "\n".join(text)
        return inst


def create(doc_path, output_path="sqlite://"):
    '''
    Parse the relational table-like XML file provided by http://www.unimod.org/downloads.html
    and convert each <tag>_row into an equivalent database entry.

    By default the table will be held in memory.
    '''
    tree = preprocess_xml(doc_path)
    engine = create_engine(output_path)
    Base.metadata.create_all(engine)
    session = sessionmaker(bind=engine, autoflush=False)()
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=sa_exc.SAWarning)
        for model in Base._decl_class_registry.values():
            if hasattr(model, "_tag_name") and hasattr(model, "from_tag"):
                for tag in tree.iterfind(".//" + model._tag_name):
                    session.add(model.from_tag(tag))
            session.commit()
    return session


def session(path="sqlite:///unimod.db"):
    engine = create_engine(path)
    Base.metadata.create_all(engine)
    session = sessionmaker(bind=engine, autoflush=False)()
    return session


class Unimod(object):
    def __init__(self, path=None):
        if path is None:
            self.path = None
            self.session = create(_unimod_xml_download_url)
        else:
            self.path = path
            try:
                self.session = session(path)
                if self.session.query(Modification).first() is None:
                    raise Exception()
            except:
                # Database may not yet exist at that location
                self.session = create(_unimod_xml_download_url, path)
                self.session.query(Modification).first()

    def get(self, identifier, strict=True):
        if isinstance(identifier, int):
            mod = self.session.query(Modification).get(identifier)
            if mod is None:
                raise KeyError(identifier)
            return mod
        elif isinstance(identifier, basestring):
            if strict:
                mod = self.session.query(Modification).filter(
                    (Modification.full_name == identifier) |
                    (Modification.code_name == identifier) |
                    (Modification.ex_code_name == identifier)).first()
                if mod is None:
                    alt_name = self.session.query(AlternativeName).filter(
                        AlternativeName.alt_name == identifier).first()
                    if alt_name is None:
                        raise KeyError(identifier)
                    mod = alt_name.modification
                return mod
            else:
                qname = "%%%s%%" % identifier
                mod = self.session.query(Modification).filter(
                    (Modification.full_name.like(qname)) |
                    (Modification.code_name.like(qname)) |
                    (Modification.ex_code_name.like(qname))).first()
                if mod is None:
                    alt_name = self.session.query(AlternativeName).filter(
                        AlternativeName.alt_name.like(qname)).first()
                    if alt_name is None:
                        raise KeyError(identifier)
                    mod = alt_name.modification
                return mod

    by_title = by_name = get

    __getitem__ = get

    @property
    def mods(self):
        return self.session.query(Modification).all()

    def __iter__(self):
        return iter(self.session.query(Modification).yield_per(1000))


def load(path=None):
    return Unimod(path)


def unimod_modification_to_dict(mod):
    spec = {
        "approved": mod.approved,
        "avge_mass": mod.average_mass,
        "composition": mod.composition,
        "full_name": mod.full_name,
        "mono_mass": mod.monoisotopic_mass,
        "record_id": mod.id,
        "specificity": [
            unimod_specificity_to_dict(s) for s in mod.specificities
        ],
        "title": mod.code_name,
        "alt_names": mod.alternative_names + [mod.ex_code_name]
    }
    return spec


def unimod_specificity_to_dict(specificity):
    spec = {
        "classification": str(specificity.classification),
        "hidden": specificity.hidden,
        "position": str(specificity.position),
        "site": specificity.amino_acid,
        "spec_group": specificity.group
    }
    spec['neutral_losses'] = [
        unimod_neutral_loss_to_dict(nl) for nl in specificity.neutral_losses
    ]
    return spec


def unimod_neutral_loss_to_dict(neutral_loss):
    spec = {
        "composition": neutral_loss.composition,
        "mono_mass": neutral_loss.monoisotopic_mass,
    }
    return spec


if __name__ == "__main__":
    import sys
    import json

    db = Unimod()
    mods = db.mods
    encoded = [unimod_modification_to_dict(mod) for mod in mods]
    encoded.sort(key=lambda x: x["record_id"])
    # need method to impose patches
    try:
        fname = sys.argv[1]
        with open(fname, 'wt') as fh:
            json.dump(encoded, fh, sort_keys=True, indent=4)
    except IndexError:
        json.dump(encoded, sys.stdout, sort_keys=True, indent=4)
