"""
CIRpy

Python interface for the Chemical Identifier Resolver (CIR) by the CADD Group at the NCI/NIH.
https://github.com/mcs07/CIRpy
"""

from __future__ import annotations

import functools
import inspect
import logging
import os
from urllib.parse import quote, urlencode
from urllib.request import urlopen

from lxml import etree

__author__ = "Matt Swain"
__email__ = "m.swain@me.com"
__version__ = "1.0.2"
__license__ = "MIT"

log = logging.getLogger("cirpy")
log.addHandler(logging.NullHandler())

API_BASE = "https://cactus.nci.nih.gov/chemical/structure"
FILE_FORMATS = {
    "alc",
    "cdxml",
    "cerius",
    "charmm",
    "cif",
    "cml",
    "ctx",
    "gjf",
    "gromacs",
    "hyperchem",
    "jme",
    "maestro",
    "mol",
    "mol2",
    "mrv",
    "pdb",
    "sdf3000",
    "sln",
    "xyz",
}


def construct_api_url(
    input,
    representation,
    resolvers=None,
    get3d=False,
    tautomers=False,
    xml=True,
    **kwargs,
):
    """Return the URL for the desired API endpoint.

    :param string input: Chemical identifier to resolve
    :param string representation: Desired output representation
    :param list(str) resolvers: (Optional) Ordered list of resolvers to use
    :param bool get3d: (Optional) Whether to return 3D coordinates (where applicable)
    :param bool tautomers: (Optional) Whether to return all tautomers
    :param bool xml: (Optional) Whether to return full XML response
    :returns: CIR API URL
    :rtype: str
    """
    # File formats require representation=file and the format in the querystring
    if representation in FILE_FORMATS:
        kwargs["format"] = representation
        representation = "file"
    # Prepend input with 'tautomers:' to return all tautomers
    if tautomers:
        input = "tautomers:%s" % input
    url = "%s/%s/%s" % (API_BASE, quote(input), representation)
    if xml:
        url += "/xml"
    if resolvers:
        kwargs["resolver"] = ",".join(resolvers)
    if get3d:
        kwargs["get3d"] = True
    if kwargs:
        url += "?%s" % urlencode(kwargs)
    return url


def request(
    input, representation, resolvers=None, get3d=False, tautomers=False, **kwargs
):
    """Make a request to CIR and return the XML response.

    :param string input: Chemical identifier to resolve
    :param string representation: Desired output representation
    :param list(string) resolvers: (Optional) Ordered list of resolvers to use
    :param bool get3d: (Optional) Whether to return 3D coordinates (where applicable)
    :param bool tautomers: (Optional) Whether to return all tautomers
    :returns: XML response from CIR
    :rtype: Element
    :raises HTTPError: if CIR returns an error code
    :raises ParseError: if CIR response is uninterpretable
    """
    url = construct_api_url(
        input, representation, resolvers, get3d, tautomers, **kwargs
    )
    log.debug("Making request: %s", url)
    response = urlopen(url)
    return etree.parse(response).getroot()


class Result:
    """A single result returned by CIR."""

    def __init__(self, input, notation, input_format, resolver, representation, value):
        """

        :param string input: Originally supplied input identifier that produced this result
        :param string notation: Identifier matched by the resolver or tautomer ID
        :param string input_format: Format of the input as interpreted by the resolver
        :param string resolver: Resolver used to produce this result
        :param string representation: Requested output representation
        :param value: Actual result value
        :type value: string or list(string)
        """
        self.input = input
        self.representation = representation
        self.resolver = resolver
        self.input_format = input_format
        self.notation = notation
        self.value = value

    def __repr__(self):
        return (
            "Result(input=%r, representation=%r, resolver=%r, input_format=%r, notation=%r, value=%r)"
            % (
                self.input,
                self.representation,
                self.resolver,
                self.input_format,
                self.notation,
                self.value,
            )
        )

    def __str__(self):
        return self.value

    def __eq__(self, other):
        return isinstance(other, type(self)) and self.__dict__ == other.__dict__

    def __getitem__(self, prop):
        """Allow dict-style access to attributes to ease transition from when results were dicts."""
        if prop in self.__dict__:
            return getattr(self, prop)
        raise KeyError(prop)

    def __setitem__(self, prop, val):
        """Allow dict-style setting of attributes to ease transition from when results were dicts."""
        setattr(self, prop, val)

    def __contains__(self, prop):
        """Allow dict-style checking of attributes to ease transition from when results were dicts."""
        return prop in self.__dict__

    def to_dict(self):
        """Return a dictionary containing Result data."""
        return self.__dict__


def query(
    input, representation, resolvers=None, get3d=False, tautomers=False, **kwargs
):
    """Get all results for resolving input to the specified output representation.

    :param string input: Chemical identifier to resolve
    :param string representation: Desired output representation
    :param list(string) resolvers: (Optional) Ordered list of resolvers to use
    :param bool get3d: (Optional) Whether to return 3D coordinates (where applicable)
    :param bool tautomers: (Optional) Whether to return all tautomers
    :returns: List of resolved results
    :rtype: list(Result)
    :raises HTTPError: if CIR returns an error code
    :raises ParseError: if CIR response is uninterpretable
    """
    tree = request(input, representation, resolvers, get3d, tautomers, **kwargs)
    results = []
    for data in tree.findall(".//data"):
        value = [item.text for item in data.findall("item")]
        result = Result(
            input=tree.attrib["string"],
            representation=tree.attrib["representation"],
            resolver=data.attrib["resolver"],
            input_format=data.attrib["string_class"],
            notation=data.attrib["notation"],
            value=value[0] if len(value) == 1 else value,
        )
        results.append(result)
    log.debug("Received %s query results", len(results))
    return results


def resolve(input, representation, resolvers=None, get3d=False, **kwargs):
    """Resolve input to the specified output representation.

    :param string input: Chemical identifier to resolve
    :param string representation: Desired output representation
    :param list(string) resolvers: (Optional) Ordered list of resolvers to use
    :param bool get3d: (Optional) Whether to return 3D coordinates (where applicable)
    :returns: Output representation or None
    :rtype: string or None
    :raises HTTPError: if CIR returns an error code
    :raises ParseError: if CIR response is uninterpretable
    """
    # Take first result from XML query
    results = query(input, representation, resolvers, False, get3d, **kwargs)
    result = results[0].value if results else None
    return result


def resolve_image(
    input,
    resolvers=None,
    fmt="png",
    width=300,
    height=300,
    frame=False,
    crop=None,
    bgcolor=None,
    atomcolor=None,
    hcolor=None,
    bondcolor=None,
    framecolor=None,
    symbolfontsize=11,
    linewidth=2,
    hsymbol="special",
    csymbol="special",
    stereolabels=False,
    stereowedges=True,
    header=None,
    footer=None,
    **kwargs,
):
    """Resolve input to a 2D image depiction.

    :param string input: Chemical identifier to resolve
    :param list(string) resolvers: (Optional) Ordered list of resolvers to use
    :param string fmt: (Optional) gif or png image format (default png)

    :param int width: (Optional) Image width in pixels (default 300)
    :param int height: (Optional) Image height in pixels (default 300)
    :param bool frame: (Optional) Whether to show border frame (default False)
    :param int crop: (Optional) Crop image with specified padding

    :param int symbolfontsize: (Optional) Atom label font size (default 11)
    :param int linewidth: (Optional) Bond line width (default 2)

    :param string bgcolor: (Optional) Background color
    :param string atomcolor: (Optional) Atom label color
    :param string hcolor: (Optional) Hydrogen atom label color
    :param string bondcolor: (Optional) Bond color
    :param string framecolor: (Optional) Border frame color

    :param bool hsymbol: (Optional) Hydrogens: all, special or none (default special)
    :param bool csymbol: (Optional) Carbons: all, special or none (default special)
    :param bool stereolabels: (Optional) Whether to show stereochemistry labels (default False)
    :param bool stereowedges: (Optional) Whether to show wedge/dash bonds (default True)
    :param string header: (Optional) Header text above structure
    :param string footer: (Optional) Footer text below structure

    """
    # Aggregate all arguments into kwargs
    frame = inspect.currentframe()
    if frame:
        args, _, _, values = inspect.getargvalues(frame)
        for arg in args:
            if values[arg] is not None:
                kwargs[arg] = values[arg]
    else:
        print("Could not get the current frame. Performing alternative action...")
    # Turn off anti-aliasing for transparent background
    if kwargs.get("bgcolor") == "transparent":
        kwargs["antialiasing"] = False
    # Renamed parameters
    if "stereolabels" in kwargs:
        kwargs["showstereo"] = kwargs.pop("stereolabels")
    if "fmt" in kwargs:
        kwargs["format"] = kwargs.pop("fmt")
    # Toggle stereo wedges
    if "stereowedges" in kwargs:
        status = kwargs.pop("stereowedges")
        kwargs.update({"wedges": status, "dashes": status})
    # Constant values
    kwargs.update({"representation": "image", "xml": False})
    url = construct_api_url(**kwargs)
    log.debug("Making image request: %s", url)
    response = urlopen(url)
    return response.read()


# TODO: Support twirl as fmt paramter?
# TODO: ipython html repr twirl, ipython png repr image


def download(
    input,
    filename,
    representation,
    overwrite=False,
    resolvers=None,
    get3d=False,
    **kwargs,
):
    """Convenience function to save a CIR response as a file.

    This is just a simple wrapper around the resolve function.

    :param string input: Chemical identifier to resolve
    :param string filename: File path to save to
    :param string representation: Desired output representation
    :param bool overwrite: (Optional) Whether to allow overwriting of an existing file
    :param list(string) resolvers: (Optional) Ordered list of resolvers to use
    :param bool get3d: (Optional) Whether to return 3D coordinates (where applicable)
    :raises HTTPError: if CIR returns an error code
    :raises ParseError: if CIR response is uninterpretable
    :raises IOError: if overwrite is False and file already exists
    """
    result = resolve(input, representation, resolvers, get3d, **kwargs)
    # Just log and return if nothing resolved
    if not result:
        log.debug("No file to download.")
        return
    # Only overwrite an existing file if explicitly instructed to.
    if not overwrite and os.path.isfile(filename):
        raise OSError(
            "%s already exists. Use 'overwrite=True' to overwrite it." % filename
        )
    # Ensure file ends with a newline
    if not result.endswith("\n"):
        result += "\n"
    with open(filename, "w") as f:
        f.write(result)


def memoized_property(fget):
    """Decorator to create memoized properties."""
    attr_name = f"_{fget.__name__}"

    @functools.wraps(fget)
    def fget_memoized(self):
        if not hasattr(self, attr_name):
            setattr(self, attr_name, fget(self))
        return getattr(self, attr_name)

    return property(fget_memoized)


class Molecule:
    """Class to hold and cache the structure information for a given CIR input."""

    def __init__(self, input, resolvers=None, get3d=False, **kwargs):
        """Initialize with a resolver input."""
        self.input = input
        self.resolvers = resolvers
        self.get3d = get3d
        self.kwargs = kwargs
        log.debug("Instantiated Molecule: %s" % self)

    def __repr__(self):
        return "Molecule(input=%r, resolvers=%r, get3d=%r, kwargs=%r)" % (
            self.input,
            self.resolvers,
            self.get3d,
            self.kwargs,
        )

    @memoized_property
    def stdinchi(self):
        """Standard InChI."""
        return resolve(self.input, "stdinchi", self.resolvers, **self.kwargs)

    @memoized_property
    def stdinchikey(self):
        """Standard InChIKey."""
        return resolve(self.input, "stdinchikey", self.resolvers, **self.kwargs)

    @memoized_property
    def inchi(self):
        """Non-standard InChI. (Uses options DONOTADDH W0 FIXEDH RECMET NEWPS SPXYZ SAsXYZ Fb Fnud)."""
        return resolve(self.input, "inchi", self.resolvers, **self.kwargs)

    @memoized_property
    def smiles(self):
        """SMILES string."""
        return resolve(self.input, "smiles", self.resolvers, **self.kwargs)

    @memoized_property
    def ficts(self):
        """FICTS NCI/CADD hashed structure identifier."""
        return resolve(self.input, "ficts", self.resolvers, **self.kwargs)

    @memoized_property
    def ficus(self):
        """FICuS NCI/CADD hashed structure identifier."""
        return resolve(self.input, "ficus", self.resolvers, **self.kwargs)

    @memoized_property
    def uuuuu(self):
        """uuuuu NCI/CADD hashed structure identifier."""
        return resolve(self.input, "uuuuu", self.resolvers, **self.kwargs)

    @memoized_property
    def hashisy(self):
        """CACTVS HASHISY identifier."""
        return resolve(self.input, "hashisy", self.resolvers, **self.kwargs)

    @memoized_property
    def sdf(self):
        """SDF file."""
        return resolve(self.input, "sdf", self.resolvers, **self.kwargs)

    @memoized_property
    def names(self):
        """List of chemical names."""
        return resolve(self.input, "names", self.resolvers, **self.kwargs)

    @memoized_property
    def iupac_name(self):
        """IUPAC approved name."""
        return resolve(self.input, "iupac_name", self.resolvers, **self.kwargs)

    @memoized_property
    def cas(self):
        """CAS registry numbers."""
        return resolve(self.input, "cas", self.resolvers, **self.kwargs)

    @memoized_property
    def mw(self):
        """Molecular weight."""
        return resolve(self.input, "mw", self.resolvers, **self.kwargs)

    @memoized_property
    def formula(self):
        """Molecular formula"""
        return resolve(self.input, "formula", self.resolvers, **self.kwargs)

    @memoized_property
    def h_bond_donor_count(self):
        """Hydrogen bond donor count."""
        return resolve(self.input, "h_bond_donor_count", self.resolvers, **self.kwargs)

    @memoized_property
    def h_bond_acceptor_count(self):
        """Hydrogen bond acceptor count."""
        return resolve(
            self.input, "h_bond_acceptor_count", self.resolvers, **self.kwargs
        )

    @memoized_property
    def h_bond_center_count(self):
        """Hydrogen bond center count."""
        return resolve(self.input, "h_bond_center_count", self.resolvers, **self.kwargs)

    @memoized_property
    def rule_of_5_violation_count(self):
        """Rule of 5 violation count."""
        return resolve(
            self.input, "rule_of_5_violation_count", self.resolvers, **self.kwargs
        )

    @memoized_property
    def rotor_count(self):
        """Rotor count."""
        return resolve(self.input, "rotor_count", self.resolvers, **self.kwargs)

    @memoized_property
    def effective_rotor_count(self):
        """Effective rotor count."""
        return resolve(
            self.input, "effective_rotor_count", self.resolvers, **self.kwargs
        )

    @memoized_property
    def ring_count(self):
        """Ring count."""
        return resolve(self.input, "ring_count", self.resolvers, **self.kwargs)

    @memoized_property
    def ringsys_count(self):
        """Ring system count."""
        return resolve(self.input, "ringsys_count", self.resolvers, **self.kwargs)

    @memoized_property
    def image(self):
        """2D image depiction."""
        return resolve_image(self.input, self.resolvers, **self.kwargs)

    @property
    def image_url(self):
        """URL of a GIF image."""
        return construct_api_url(
            self.input, "image", self.resolvers, False, self.get3d, False, **self.kwargs
        )

    @property
    def twirl_url(self):
        """Url of a TwirlyMol 3D viewer."""
        return construct_api_url(
            self.input, "twirl", self.resolvers, False, self.get3d, False, **self.kwargs
        )

    def download(self, filename, representation, overwrite=False):
        """Download the resolved structure as a file.

        :param string filename: File path to save to
        :param string representation: Desired output representation
        :param bool overwrite: (Optional) Whether to allow overwriting of an existing file
        """
        download(
            self.input,
            filename,
            representation,
            overwrite,
            self.resolvers,
            self.get3d,
            **self.kwargs,
        )
