# Copyright 2020 Open Reaction Database Project Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Parses CML from the NRD.

Data is at:
https://figshare.com/articles/dataset/Chemical_reactions_from_US_patents_1976-Sep2016_/5104873

See also:
https://depth-first.com/articles/2019/01/28/the-nextmove-patent-reaction-dataset/

Usage:
    parse_uspto.py --input_pattern=<str> --name=<str> --output=<str> [--n_jobs=<int>]

Options:
    --input_pattern=<str>       Input pattern for CML files
    --name=<str>                Dataset name
    --output=<str>              Output Dataset filename
    --n_jobs=<str>              Number of parallel workers [default: 1]
"""
# pylint: disable=import-error
# pylint: disable=too-many-branches

import datetime
import glob
import logging
import os
import re
from typing import Dict, List, Tuple, Union
from xml.etree import ElementTree

import docopt
import joblib
from rdkit import RDLogger

import ord_schema
from ord_schema import message_helpers
from ord_schema import units
from ord_schema import validations
from ord_schema.proto import dataset_pb2
from ord_schema.proto import reaction_pb2

logger = logging.getLogger()

RDLogger.DisableLog("rdApp.*")  # Disable RDKit logging.

# XML namespaces.
NAMESPACES = {
    "cml": "http://www.xml-cml.org/schema",
    "cmlDict": "http://www.xml-cml.org/dictionary/cml/",
    "dl": "http://bitbucket.org/dan2097",
    "nameDict": "http://www.xml-cml.org/dictionary/cml/name/",
    "unit": "http://www.xml-cml.org/unit/",
}

UNIT_RESOLVER = units.UnitResolver()

REACTION_ROLES = {
    "catalyst": reaction_pb2.ReactionRole.CATALYST,
    "product": reaction_pb2.ReactionRole.PRODUCT,
    "reactant": reaction_pb2.ReactionRole.REACTANT,
    "solvent": reaction_pb2.ReactionRole.SOLVENT,
}

PRODUCT_STATES = {
    "crystal": reaction_pb2.ProductCompound.Texture.CRYSTAL,
    "crystalline-solid": reaction_pb2.ProductCompound.Texture.CRYSTAL,
    "crystals": reaction_pb2.ProductCompound.Texture.CRYSTAL,
    "foam": reaction_pb2.ProductCompound.Texture.FOAM,
    "foam-like": reaction_pb2.ProductCompound.Texture.FOAM,
    "foam/very": reaction_pb2.ProductCompound.Texture.FOAM,
    "needle": reaction_pb2.ProductCompound.Texture.CRYSTAL,
    "needle-like": reaction_pb2.ProductCompound.Texture.CRYSTAL,
    "needles": reaction_pb2.ProductCompound.Texture.CRYSTAL,
    "oil": reaction_pb2.ProductCompound.Texture.OIL,
    "oil-like": reaction_pb2.ProductCompound.Texture.OIL,
    "oils": reaction_pb2.ProductCompound.Texture.OIL,
    "powder": reaction_pb2.ProductCompound.Texture.POWDER,
    "powders": reaction_pb2.ProductCompound.Texture.POWDER,
    "prisms": reaction_pb2.ProductCompound.Texture.CRYSTAL,
    "semisolid": reaction_pb2.ProductCompound.Texture.SEMI_SOLID,
    "semi-solid": reaction_pb2.ProductCompound.Texture.SEMI_SOLID,
}

WORKUP_TYPES = {
    "Add": reaction_pb2.ReactionWorkup.ADDITION,
    "Cool": reaction_pb2.ReactionWorkup.TEMPERATURE,
    "Concentrate": reaction_pb2.ReactionWorkup.CONCENTRATION,
    "Dissolve": reaction_pb2.ReactionWorkup.DISSOLUTION,
    "Distill": reaction_pb2.ReactionWorkup.DISTILLATION,
    "Dry with material": reaction_pb2.ReactionWorkup.DRY_WITH_MATERIAL,
    "Extract": reaction_pb2.ReactionWorkup.EXTRACTION,
    "Filter": reaction_pb2.ReactionWorkup.FILTRATION,
    "Heat": reaction_pb2.ReactionWorkup.TEMPERATURE,
    "Sample": reaction_pb2.ReactionWorkup.ALIQUOT,
    "Stir": reaction_pb2.ReactionWorkup.STIRRING,
    "Wait": reaction_pb2.ReactionWorkup.WAIT,
    "Wash": reaction_pb2.ReactionWorkup.WASH,
}

UNIT_REPLACEMENTS = {
    "˜": " ",
    "≈": " ",
    "about ": " ",
    "approx ": " ",
    "approx. ": " ",
    "approximately ": " ",
    "around ": " ",
    "ca. ": " ",
    "yield ": " ",
    "yield: ": " ",
    "Yield ": " ",
    "Yield: ": " ",
    "1/2 ": " 0.5 ",
    "½ ": " 0.5 ",
    "half ": " 0.5 ",
    " to ": "-",
    "×10": "e",
    "−": "-",
    # Digits.
    "zero ": "0 ",
    "one ": "1 ",
    "two ": "2 ",
    "three ": "3 ",
    "four ": "4 ",
    "five ": "5 ",
    "six ": "6 ",
    "seven ": "7 ",
    "eight ": "8 ",
    "nine ": "9 ",
}


def get_tag(element: ElementTree.Element) -> str:
    """Fetches the namespace-translated element tag."""
    for key, value in NAMESPACES.items():
        if value in element.tag:
            return element.tag.replace(f"{{{value}}}", f"{key}:")
    return element.tag


def get_molecule_id(root: ElementTree.Element) -> str:
    """Fetches the ID of the molecule under this element."""
    molecules = root.findall("cml:molecule", namespaces=NAMESPACES)
    if len(molecules) != 1:
        raise NotImplementedError(len(molecules))
    return molecules[0].attrib["id"]


def resolve_units(value: str) -> ord_schema.UnitMessage:
    """Resolves a value/unit string."""
    for key, replacement in UNIT_REPLACEMENTS.items():
        value = value.replace(key, replacement)
    return UNIT_RESOLVER.resolve(value, allow_range=True)


def get_component_map(root: ElementTree.Element) -> Dict[str, str]:
    """Builds a mapping of components to inputs."""
    reaction_inputs = {}
    # Start with a separate input for each component.
    for component in root.find("cml:reactantList", namespaces=NAMESPACES):
        molecule_id = get_molecule_id(component)
        reaction_inputs[molecule_id] = molecule_id
    for component in root.find("cml:spectatorList", namespaces=NAMESPACES):
        molecule_id = get_molecule_id(component)
        reaction_inputs[molecule_id] = molecule_id
    for action in root.find("dl:reactionActionList", namespaces=NAMESPACES):
        components = []
        for component in action.findall("dl:chemical", namespaces=NAMESPACES):
            if "ref" in component.attrib:
                components.append(component.attrib["ref"])
        if components:
            key = "_".join(components)
            for component in components:
                reaction_inputs[component] = key
    return reaction_inputs


def parse_reaction(root: ElementTree.Element) -> reaction_pb2.Reaction:
    """Parses reaction CML into a Reaction message."""
    reaction = reaction_pb2.Reaction()
    component_map = get_component_map(root)
    for child in root:
        tag = get_tag(child)
        if tag == "dl:source":
            parse_source(child, reaction)
        elif tag == "dl:reactionSmiles":
            reaction.identifiers.add(type="REACTION_CXSMILES", value=child.text, is_mapped=True)
        elif tag == "cml:productList":
            outcome = reaction.outcomes.add()  # Add a single outcome.
            for product in child:
                parse_product(product, outcome.products.add())
        elif tag in ["cml:reactantList", "cml:spectatorList"]:
            for compound in child:
                molecule_id = get_molecule_id(compound)
                reaction_input = reaction.inputs[component_map[molecule_id]]
                parse_reactant(compound, reaction_input.components.add())
        elif tag == "dl:reactionActionList":
            for action in child:
                parse_conditions(action, reaction)
                parse_workup(action, reaction)
        else:
            raise NotImplementedError(child)
    # We surely didn't capture everything.
    reaction.conditions.conditions_are_dynamic = True
    reaction.conditions.details = "See reaction.notes.procedure_details."
    return reaction


def parse_source(root: ElementTree.Element, reaction: reaction_pb2.Reaction):
    """Adds provenance information to a Reaction."""
    for child in root:
        tag = get_tag(child)
        if tag == "dl:documentId":
            reaction.provenance.patent = child.text
        elif tag == "dl:paragraphText":
            reaction.notes.procedure_details = child.text
        elif tag in ["dl:headingText", "dl:paragraphNum"]:
            continue  # Ignored.
        else:
            raise NotImplementedError(child)


def parse_product(root: ElementTree.Element, product_compound: reaction_pb2.ProductCompound):
    """Adds product information to a ProductCompound."""
    role = root.attrib.get("role")
    if role:
        product_compound.reaction_role = REACTION_ROLES[role]
    for child in root:
        tag = get_tag(child)
        if tag == "dl:entityType":
            pass
        elif tag == "cml:molecule":
            parse_molecule(child, product_compound)
        elif tag == "cml:amount":
            parse_product_amount(child, product_compound)
        elif tag == "cml:identifier":
            parse_identifier(child, product_compound)
        elif tag == "dl:state":
            texture = PRODUCT_STATES.get(child.text.lower(), reaction_pb2.ProductCompound.Texture.CUSTOM)
            product_compound.texture.type = texture
            product_compound.texture.details = child.text
        elif tag == "dl:appearance":
            product_compound.isolated_color = child.text
        else:
            raise NotImplementedError(child)


def parse_molecule(
    root: ElementTree.Element,
    compound: Union[reaction_pb2.Compound, reaction_pb2.ProductCompound],
):
    """Adds NAME identifiers to a Compound."""
    for child in root:
        tag = get_tag(child)
        if tag in ["cml:name", "dl:nameResolved"]:
            compound.identifiers.add(type="NAME", value=child.text)
        else:
            raise NotImplementedError(child)


def parse_product_amount(root: ElementTree.Element, product_compound: reaction_pb2.ProductCompound):
    """Adds amount information to a ProductCompound."""
    property_type = root.attrib[f'{{{NAMESPACES["dl"]}}}propertyType']
    if "PERCENTYIELD" in property_type:
        match = re.search(r"(\d+\.?\d*)", root.text)
        if not match:
            logger.debug(f"AMOUNT: {root.text}")
            return
        measurement = product_compound.measurements.add()
        measurement.type = measurement.YIELD
        measurement.percentage.value = float(match.group(1))
    else:
        measurement = product_compound.measurements.add()
        measurement.type = measurement.AMOUNT
        parse_amount(root, measurement)
    measurement.details = property_type


def parse_identifier(
    root: ElementTree.Element,
    compound: Union[reaction_pb2.Compound, reaction_pb2.ProductCompound],
):
    """Adds a SMILES or INCHI identifier to a Compound."""
    kind = root.attrib["dictRef"]
    value = root.attrib["value"]
    if kind == "cml:smiles":
        compound.identifiers.add(type="SMILES", value=value)
    elif kind == "cml:inchi":
        compound.identifiers.add(type="INCHI", value=value)
    else:
        raise NotImplementedError(kind)


def parse_reactant(root: ElementTree.Element, compound: reaction_pb2.Compound):
    """Populates an input Compound."""
    role = root.attrib.get("role")
    if role:
        compound.reaction_role = REACTION_ROLES[role]
    for child in root:
        tag = get_tag(child)
        if tag == "dl:entityType":
            if child.text not in ["exact", "chemicalClass", "definiteReference"]:
                raise NotImplementedError(child.text)
        elif tag == "cml:molecule":
            parse_molecule(child, compound)
        elif tag == "cml:amount":
            parse_amount(child, compound)
        elif tag == "cml:identifier":
            parse_identifier(child, compound)
        elif tag in ["dl:state", "dl:appearance"]:
            continue  # Not supported by Compound.
        else:
            raise NotImplementedError(child)


def parse_amount(
    root: ElementTree.Element,
    compound: Union[reaction_pb2.Compound, reaction_pb2.ProductMeasurement],
):
    """Parses an amount."""
    property_type = root.attrib[f'{{{NAMESPACES["dl"]}}}propertyType']
    if property_type in ["MOLARITY", "PH"]:
        return
    if compound.amount.WhichOneof("kind") not in ["mass", "volume"]:
        # Don't overwrite existing Mass or Volume with Moles.
        try:
            amount = resolve_units(root.text)
        except (KeyError, ValueError) as error:
            logger.debug(f'UNITS: {error} ("{root.text}")')
            return
        if isinstance(amount, reaction_pb2.Mass):
            compound.amount.mass.CopyFrom(amount)
        elif isinstance(amount, reaction_pb2.Moles):
            compound.amount.moles.CopyFrom(amount)
        elif isinstance(amount, reaction_pb2.Volume):
            compound.amount.volume.CopyFrom(amount)
        else:
            raise NotImplementedError(amount)


def parse_conditions(root: ElementTree.Element, reaction: reaction_pb2.Reaction):
    """Parses reaction conditions."""
    del reaction  # Unused.
    if not root.findall("cml:chemical", namespaces=NAMESPACES):
        return  # Refers to an added component; is a workup.
    action = root.attrib["action"]
    del action  # Unused.
    return  # TODO(kearnes): Implement this?


def parse_workup(root: ElementTree.Element, reaction: reaction_pb2.Reaction):
    """Parses a workup step."""
    if root.findall("dl:chemical", namespaces=NAMESPACES):
        return  # Refers to an input component; not a workup.
    action = root.attrib["action"]
    components = []
    for component in root.findall("cml:chemical", namespaces=NAMESPACES):
        entity_type = component.find("dl:entityType", namespaces=NAMESPACES)
        assert entity_type is not None  # Type hint.
        if entity_type.text in ["exact", "chemicalClass", "definiteReference"]:
            components.append(component)
    if action == "Dry" and components:
        action = "Dry with material"
    details = root.find("dl:phraseText", namespaces=NAMESPACES)
    assert details is not None  # Type hint.
    details = details.text
    if action in ["Purify", "Recover"] and details:
        # Make some actions more specific.
        # TODO(kearnes): This could be expanded.
        if "distill" in details.lower():
            action = "Distill"
        elif "filtration" in details.lower():
            action = "Filter"
    workup_type = WORKUP_TYPES.get(action, reaction_pb2.ReactionWorkup.CUSTOM)
    if workup_type == reaction_pb2.ReactionWorkup.CUSTOM:
        logger.debug(f'CUSTOM workup "{action}": {details}')
    workup = reaction.workups.add(type=workup_type)
    for component in components:
        compound = workup.input.components.add()
        parse_reactant(component, compound)
        compound.reaction_role = reaction_pb2.ReactionRole.WORKUP
    for child in root:
        tag = get_tag(child)
        if tag == "cml:chemical":
            pass  # Components are handled all together above.
        elif tag == "dl:phraseText":
            workup.details = child.text
        elif tag == "dl:parameter":
            try:
                parse_parameter(child, workup)
            except ValueError as error:
                logger.debug(f"PARAMETER: {ElementTree.tostring(child)}: {error}")
        elif tag in ["dl:atmosphere"]:
            pass  # Not supported by ReactionWorkup.
        else:
            raise NotImplementedError(child)


def parse_parameter(root: ElementTree.Element, workup: reaction_pb2.ReactionWorkup):
    """Parses a workup value."""
    kind = root.attrib["propertyType"]
    if kind == "Time":
        if root.text in ["overnight", "Overnight"]:
            workup.duration.value = 8
            workup.duration.precision = 8
            workup.duration.units = reaction_pb2.Time.HOUR
        else:
            try:
                workup.duration.CopyFrom(resolve_units(root.text))
            except (KeyError, ValueError) as error:
                logger.debug(f'UNITS: {error} ("{root.text}")')
    elif kind == "Temperature":
        if root.text.lower() in [
            "room temperature",
            "room temp",
            "rt",
            "r.t",
            "r.t.",
            "ambient temperature",
            "ambient temp",
        ]:
            workup.temperature.control.type = reaction_pb2.TemperatureConditions.TemperatureControl.AMBIENT
        else:
            value = root.text.rstrip(".").replace("° ", "°")
            try:
                temperature = resolve_units(value)
                if temperature.units == temperature.CELSIUS and temperature.value < -274:
                    raise ValueError("bad temperature")
                if temperature.precision < 0:
                    raise ValueError("bad temperature")
                workup.temperature.setpoint.CopyFrom(temperature)
            except (KeyError, ValueError) as error:
                logger.debug(f'UNITS: {error} ("{root.text}")')
    elif kind in ["Frequency", "Length", "Pressure"]:
        return  # Not supported in ReactionWorkup.
    else:
        raise NotImplementedError(kind)


def clean_reaction(reaction: reaction_pb2.Reaction):
    """Cleans a reaction so it will pass validations."""
    # Add a placeholder amount to components with no amount information.
    empty_amount = reaction_pb2.Moles(value=0, precision=1, units="MOLE")
    components = message_helpers.find_submessages(reaction, reaction_pb2.Compound)
    for component in components:
        if not component.amount.WhichOneof("kind"):
            component.amount.moles.CopyFrom(empty_amount)
    # Adjust identifier types as needed.
    identifiers = message_helpers.find_submessages(reaction, reaction_pb2.CompoundIdentifier)
    for identifier in identifiers:
        output = validations.validate_message(identifier, raise_on_error=False)
        if output.errors:
            old_type = reaction_pb2.CompoundIdentifier.IdentifierType.Name(identifier.type)
            identifier.details = f"Originally defined as {old_type}"
            identifier.type = reaction_pb2.CompoundIdentifier.CUSTOM
    # Adjust workup types as needed.
    for workup in reaction.workups:
        output = validations.validate_message(workup, raise_on_error=False)
        if output.errors:
            workup.type = reaction_pb2.ReactionWorkup.CUSTOM
    # Move some workups into ReactionConditions/ReactionOutcome.
    workups = [workup for workup in reaction.workups]  # pylint: disable=unnecessary-comprehension
    del reaction.workups[:]
    temperature_conditions = False
    stirring_conditions = False
    reaction_time = False
    for workup in workups:
        if workup.type == workup.STIRRING and not stirring_conditions:
            reaction.conditions.stirring.CopyFrom(workup.stirring)
            reaction.conditions.stirring.method.type = reaction_pb2.StirringConditions.StirringMethod.CUSTOM
            reaction.conditions.stirring.method.details = workup.details
            stirring_conditions = True
            if workup.HasField("temperature") and not temperature_conditions:
                reaction.conditions.temperature.CopyFrom(workup.temperature)
                temperature_conditions = True
            if workup.HasField("duration") and not reaction_time:
                reaction.outcomes[0].reaction_time.CopyFrom(workup.duration)
                reaction_time = True
        elif workup.type == workup.WAIT and workup.HasField("duration") and not reaction_time:
            reaction.outcomes[0].reaction_time.CopyFrom(workup.duration)
            reaction_time = True
        elif workup.type == workup.TEMPERATURE and workup.HasField("temperature") and not temperature_conditions:
            reaction.conditions.temperature.CopyFrom(workup.temperature)
            temperature_conditions = True
        else:
            reaction.workups.add().CopyFrom(workup)


def run(filename: str) -> Tuple[List[reaction_pb2.Reaction], List[reaction_pb2.Reaction]]:
    """Parses reactions from a single CML file."""
    RDLogger.DisableLog("rdApp.*")  # Disable RDKit logging.
    tree = ElementTree.parse(filename)
    root = tree.getroot()
    reactions = []
    failures = []
    for reaction_cml in root.iterfind("cml:reaction", namespaces=NAMESPACES):
        try:
            reaction = parse_reaction(reaction_cml)
        except (KeyError, NotImplementedError) as error:
            raise ValueError(ElementTree.dump(reaction_cml)) from error
        event = reaction_pb2.RecordEvent(
            time=dict(value=str(datetime.datetime.now())),
            person=dict(
                username="skearnes",
                name="Steven Kearnes",
                orcid="0000-0003-4579-4388",
                organization="Google LLC",
                email="kearnes@google.com",
            ),
        )
        reaction.provenance.record_created.CopyFrom(event)
        reaction.provenance.doi = "10.6084/m9.figshare.5104873.v1"
        clean_reaction(reaction)
        # Check reaction SMILES.
        try:
            message_helpers.validate_reaction_smiles(reaction.identifiers[0].value.split()[0])
        except ValueError:
            failures.append(reaction)
            continue
        output = validations.validate_message(reaction, raise_on_error=False)
        if output.errors:
            message = "\n".join(output.errors)
            raise validations.ValidationError(f"{reaction}\n{message}")
        reactions.append(reaction)
    return reactions, failures


def main(kwargs):
    filenames = sorted(glob.glob(kwargs["--input_pattern"]))
    all_reactions = joblib.Parallel(n_jobs=kwargs["--n_jobs"], verbose=True)(
        joblib.delayed(run)(filename) for filename in filenames
    )
    reactions = []
    failures = []
    for file_reactions, file_failures in all_reactions:
        reactions.extend(file_reactions)
        failures.extend(file_failures)
    dataset = dataset_pb2.Dataset(reactions=reactions, name=kwargs["--name"])
    basenames = [os.path.basename(filename) for filename in filenames]
    dataset.description = f'CML filenames: {",".join(basenames)}'
    if kwargs["--output"] and reactions:
        message_helpers.write_message(dataset, kwargs["--output"])
        if failures:
            failure_dataset = dataset_pb2.Dataset(reactions=failures, name=kwargs["--name"])
            message_helpers.write_message(failure_dataset, kwargs["--output"] + ".failures.pb")


if __name__ == "__main__":
    main(docopt.docopt(__doc__))
