# Copyright 2020 University of Groningen
#
# 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.

"""
High level API for the polyply itp generator
"""
import sys
import networkx as nx
from pathlib import Path
import vermouth
import vermouth.forcefield
from vermouth.log_helpers import StyleAdapter, get_logger
# patch to get rid of martinize dependency
try:
    from vermouth.file_writer import deferred_open
except ImportError:
    from vermouth.file_writer import open
    deferred_open = open
from vermouth.file_writer import DeferredFileWriter
from vermouth.citation_parser import citation_formatter
from vermouth.graph_utils import make_residue_graph
from polyply import (MetaMolecule, ApplyLinks, MapToMolecule)
from polyply.src.graph_utils import find_missing_edges
from .load_library import load_ff_library
from .gen_dna import complement_dsDNA
from .simple_seq_parsers import parse_simple_seq_string
from .apply_modifications import ApplyModifications

LOGGER = StyleAdapter(get_logger(__name__))

def gen_params(name="polymer", outpath=Path("polymer.itp"), inpath=[],
               lib=None, seq=None, seq_file=None,
               dsdna=False, mods=[], protter=False):
    """
    The parameter file in polyply is generated by matching residue
    parameters (defined as vermouth.Block) to a residue graph of a
    molecule. All interactions arising from linking two residues are
    matched when there is an edge in a residue graph. These interactions
    are defined as vermouth.Links. Blocks and Links can be defined in one
    or more vermouth style .ff files provided by `inpath`. Alternatively,
    one can resort to the definitions stored in a particular library, which
    is chosen by providing the `lib` option. Both can be provided at the
    same time whereby the file parameters overwrite any equivalent parameters
    in the library files.

    There is no restriction with regard to how a molecule is partitioned
    into residues. For a protein the primary amino acid sequence may be
    provided and for synthetic polymers their monomer sequence or graph
    in case of a branched molecule. However, it is common that the terminal
    residues in a protein are different (e.g. protonated, or acetylated) from
    the other amino acids but no special residue name is provided. To account
    for this situation the `protter` flag can be used to specify how an
    aminoacid at the terminal of a protein is different from the canonical
    aminoacid found in the chain.

    DNA presents another special case. For DNA the nucleobase sequence of one
    strand may be provided. When `dsdna` is `True` the complementary
    sequence is inferred by assuming canonical base pairing. Parameters for the
    complementary strands are also written to the itp file.

    The residue graph can be provided in several ways:
    (1) A list of strings that specify monomer and number of repeat units
    (e.g. PS:5 PEO:5) for a block copolymer of PS and PEO. This route only
    allows for linear molecules.
    (2) a single CGsmiles string may be provided that describes the entire
    molecule. Matching is done on the finest level of the CGsmiles string.
    CGsmiles strings also allow definition of cyclic or branched molecules.
    (3) a .json file in networkx node-link data format
    (4) a .json file of a linear squence where each residue is seperated by
    a space from the preceding one. For a protein `ALA TRP GLU` or for a
    synthetic polymer `PS PS PS PS PS PEO PEO PEO PEO PEO`.

    Parameters
    ----------
    name: str
        name of the molecule in the itp file
    outpath: :class:`pathlib.Path`
        file path for output file
    inpath: list[pathlib.Path]
        list of paths to files with input definitions
    library: str
        name of the force field library to use
    seq: list[str]
        list of strings with format "resname:#monomers" or
        a list with one string describing the molecule in
        CGsmiles format. An abitrary amount of higher resolutions
        are permitted when using the CGsmiles string but only
        the final level is used for itp file generation.
    seqf: :class:`pathlib.Path`
        file path to valid sequence file (.json/.fasta/.ig/.txt)
    dsdna: bool
        find complementing DNA strand and write to itp
    mods: list[str]
        list of modifications to apply
    protter: bool
        apply protein terminal modifications
    """
    # Import of Itp and FF files
    LOGGER.info("reading input and library files",  type="step")
    force_field = load_ff_library(name, lib, inpath)

    # Generate the MetaMolecule
    if seq:
        LOGGER.info("reading sequence from command",  type="step")
        # We are dealing with a cgsmiles string
        if len(seq) == 1 and seq[0].startswith("{"):
            meta_molecule = MetaMolecule.from_cgsmiles_str(cgsmiles_str=seq[0],
                                                           force_field=force_field,
                                                           mol_name=name)
        else:
            monomers = parse_simple_seq_string(seq)
            meta_molecule = MetaMolecule.from_monomer_seq_linear(monomers=monomers,
                                                                 force_field=force_field,
                                                                 mol_name=name)
    elif seq_file:
        LOGGER.info("reading sequence from file",  type="step")
        meta_molecule = MetaMolecule.from_sequence_file(force_field, seq_file, name)

    # Generate complementary DNA strand
    if dsdna:
        complement_dsDNA(meta_molecule)

    # Do transformationa and apply link
    LOGGER.info("mapping sequence to molecule",  type="step")
    meta_molecule = MapToMolecule(force_field).run_molecule(meta_molecule)
    LOGGER.info("applying links between residues",  type="step")
    meta_molecule = ApplyLinks().run_molecule(meta_molecule)

    meta_molecule = ApplyModifications(modifications=mods,
                                       meta_molecule=meta_molecule).run_molecule(meta_molecule)

    # Raise warning if molecule is disconnected
    msg = "Missing a link between residue {idxA} {resA} and residue {idxB} {resB}."
    for missing in find_missing_edges(meta_molecule, meta_molecule.molecule):
        LOGGER.warning(msg, **missing)

    with deferred_open(outpath, 'w') as outfile:
        header = [ ' '.join(sys.argv) + "\n" ]
        header.append("Please cite the following papers:")
        citation_map = meta_molecule.molecule.force_field.citations
        citation_map.update(vermouth.data.COMMON_CITATIONS)
        for citation in meta_molecule.molecule.citations:
            cite_string =  citation_formatter(citation_map[citation])
            LOGGER.info("Please cite: " + cite_string)
            header.append(cite_string)

        vermouth.gmx.itp.write_molecule_itp(meta_molecule.molecule, outfile,
                                            moltype=name, header=header)
    DeferredFileWriter().write()

    # Print molecule Log messages
    if meta_molecule.molecule.log_entries:
        print("")
    for loglevel, entries in meta_molecule.molecule.log_entries.items():
        for entry, fmt_args in entries.items():
            for fmt_arg in fmt_args:
                fmt_arg = {str(k): meta_molecule.molecule.nodes[v] for k, v in fmt_arg.items()}
                LOGGER.log(loglevel, entry, **fmt_arg, type='model')

# ducktape for renaming the itp tool
gen_itp = gen_params
