#!python
# -*- coding: utf-8 -*-
# Copyright 2018 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 Martinize2
"""

import argparse
import functools
import logging
import itertools
import textwrap
from pathlib import Path
import sys

import vermouth
import vermouth.forcefield
from vermouth.file_writer import open, DeferredFileWriter
from vermouth import DATA_PATH
from vermouth.dssp import dssp
from vermouth.dssp.dssp import (
    AnnotateDSSP,
    AnnotateMartiniSecondaryStructures,
    AnnotateResidues,
)
from vermouth.log_helpers import (StyleAdapter, BipolarFormatter,
                                  CountingHandler, TypeAdapter,
                                  ignore_warnings_and_count,)
from vermouth import selectors
from vermouth.map_input import (
    read_mapping_directory,
    generate_all_self_mappings,
    combine_mappings
)
from vermouth.citation_parser import citation_formatter

# TODO Since vermouth's __init__.py does some logging (KDTree), this may or may
# not work as intended. Investigation required.

LOGGER = TypeAdapter(logging.getLogger('vermouth'))

PRETTY_FORMATTER = logging.Formatter(fmt='{levelname:>8} - {type} - {message}',
                                     style='{')
DETAILED_FORMATTER = logging.Formatter(fmt='{levelname:>8} - {type} - {name} - {message}',
                                       style='{')

COUNTER = CountingHandler()

# Control above what level message we want to count
COUNTER.setLevel(logging.WARNING)

CONSOLE_HANDLER = logging.StreamHandler()
FORMATTER = BipolarFormatter(DETAILED_FORMATTER,
                             PRETTY_FORMATTER,
                             logging.DEBUG,
                             logger=LOGGER)
CONSOLE_HANDLER.setFormatter(FORMATTER)
LOGGER.addHandler(CONSOLE_HANDLER)
LOGGER.addHandler(COUNTER)

LOGGER = StyleAdapter(LOGGER)

VERSION = 'martinize with vermouth {}'.format(vermouth.__version__)


def read_system(path, ignore_resnames=(), ignh=None, modelidx=None):
    """
    Read a system from a PDB or GRO file.

    This function guesses the file type based on the file extension.

    The resulting system does not have a force field and may not have edges.
    """
    system = vermouth.System()
    file_extension = path.suffix.upper()[1:]  # We do not keep the dot
    if file_extension in ['PDB', 'ENT']:
        vermouth.PDBInput(str(path), exclude=ignore_resnames, ignh=ignh,
                          modelidx=modelidx).run_system(system)
    elif file_extension in ['GRO']:
        vermouth.GROInput(str(path), exclude=ignore_resnames, ignh=ignh).run_system(system)
    else:
        raise ValueError('Unknown file extension "{}".'.format(file_extension))
    return system


def pdb_to_universal(system, delete_unknown=False, force_field=None,
                     modifications=None, mutations=None,
                     bonds_from_name=True, bonds_from_dist=True, bonds_fudge=1,
                     write_graph=None, write_repair=None, write_canon=None):
    """
    Convert a system read from the PDB to a clean canonical atomistic system.
    """
    if force_field is None:
        force_field = vermouth.forcefield.get_native_force_field('universal')
    if modifications is None:
        modifications = []
    if mutations is None:
        mutations = []
    canonicalized = system.copy()
    canonicalized.force_field = force_field

    LOGGER.info('Guessing the bonds.', type='step')
    vermouth.MakeBonds(allow_name=bonds_from_name,
                       allow_dist=bonds_from_dist,
                       fudge=bonds_fudge).run_system(canonicalized)
    vermouth.MergeNucleicStrands().run_system(canonicalized)
    if write_graph is not None:
        vermouth.pdb.write_pdb(canonicalized, str(write_graph), omit_charges=True)
        DeferredFileWriter().write()

    LOGGER.debug('Annotating required mutations and modifications.', type='step')
    vermouth.AnnotateMutMod(modifications, mutations).run_system(canonicalized)
    LOGGER.info('Repairing the graph.', type='step')
    vermouth.RepairGraph(delete_unknown=delete_unknown, include_graph=False).run_system(canonicalized)
    if write_repair is not None:
        vermouth.pdb.write_pdb(canonicalized, str(write_repair),
                               omit_charges=True, nan_missing_pos=True)
        DeferredFileWriter().write()
    LOGGER.info('Dealing with modifications.', type='step')
    vermouth.CanonicalizeModifications().run_system(canonicalized)
    if write_canon is not None:
        vermouth.pdb.write_pdb(canonicalized, str(write_canon),
                               omit_charges=True, nan_missing_pos=True)
        DeferredFileWriter().write()
    vermouth.AttachMass(attribute='mass').run_system(canonicalized)
    vermouth.SortMoleculeAtoms().run_system(canonicalized) # was system
    return canonicalized


def martinize(system, mappings, to_ff, delete_unknown=False):
    """
    Convert a system from one force field to an other at lower resolution.
    """
    LOGGER.info('Creating the graph at the target resolution.', type='step')
    vermouth.DoMapping(mappings=mappings,
                       to_ff=to_ff,
                       delete_unknown=delete_unknown,
                       attribute_keep=('cgsecstruct', 'resid', 'chain'),
                       attribute_must=('resname')).run_system(system)
    LOGGER.info('Averaging the coordinates.', type='step')
    vermouth.DoAverageBead(ignore_missing_graphs=True).run_system(system)
    LOGGER.info('Applying the links.', type='step')
    vermouth.DoLinks().run_system(system)
    LOGGER.info('Placing the charge dummies.', type='step')
    vermouth.LocateChargeDummies().run_system(system)
    return system


def write_gmx_topology(system, top_path, defines=(), header=()):
    """
    Writes a Gromacs .top file for the specified system.
    """
    if not system.molecules:
        raise ValueError('No molecule in the system. Nothing to write.')

    # Write the ITP files for the molecule types, and prepare writing the
    # [ molecules ] section of the top file.
    # * We write one ITP file for each different moltype in the system, the
    #   moltype being defined by the name provided under the "moltype" meta of
    #   the molecules. If more than one molecule share the same moltype, we use
    #   the first one to write the ITP file.
    moltype_written = set()
    # * We keep track of the length of the longer moltype name, to align the
    #   [ molecules ] section.
    max_name_length = 0
    # * We keep track of groups of successive molecules with the same moltypes.
    moltype_count = []  # items will be [moltype, number of molecules]

    # Iterate over groups of successive molecules with the same moltypes. We
    # shall *NOT* sort the molecules before hand, as groups of successive
    # molecules with the same moltype can be interupted by other moltypes, and
    # we want to reflect these interuptions in the [ molecules ] section of the
    # top file.
    molecule_groups = itertools.groupby(system.molecules, key=lambda x: x.meta['moltype'])
    for moltype, molecules in molecule_groups:
        molecule = next(molecules)
        if moltype not in moltype_written:
            # A given moltype can appear more than once in the sequence of
            # molecules, without being uninterupted by other moltypes. Even in
            # that case, we want to write the ITP only once.
            with open('{}.itp'.format(moltype), 'w') as outfile:
                # here we format and merge all citations
                header[-1] = header[-1]+"\n"
                header.append("Pleas cite the following papers:")
                for citation in molecule.citations:
                    cite_string =  citation_formatter(molecule.force_field.citations[citation])
                    LOGGER.info("Please cite: " + cite_string)
                    header.append(cite_string)
                vermouth.gmx.itp.write_molecule_itp(molecule, outfile, header=header)
            this_moltype_len = len(molecule.meta['moltype'])
            if this_moltype_len > max_name_length:
                max_name_length = this_moltype_len
            moltype_written.add(moltype)
        # We already removed one element from the "molecules" generator, do not
        # forget to count it in the number of molecules in that group.
        moltype_count.append([moltype, 1 + len(list(molecules))])

    # Write the top file
    template = textwrap.dedent("""\
        {defines}
        #include "martini.itp"
        {includes}

        [ system ]
        Title of the system

        [ molecules ]
        {molecules}
    """)
    include_string = '\n'.join(
        '#include "{}.itp"'.format(molecule_type)
        for molecule_type, _ in moltype_count
    )
    molecule_string = '\n'.join(
        '{mtype:<{length}}    {num}'
        .format(mtype=mtype, num=num, length=max_name_length)
        for mtype, num in moltype_count
    )
    define_string = '\n'.join(
        '#define {}'.format(define) for define in defines
    )
    with open(str(top_path), 'w') as outfile:
        outfile.write(
            textwrap.dedent(
                template.format(
                    includes=include_string,
                    molecules=molecule_string,
                    defines=define_string,
                )
            )
        )


def _cys_argument(value):
    """
    Convert and validate the value of the cys option for argparse.

    Parameters
    ----------
    value: str
        The value given to the command line.

    Return
    ------
    str or float
        A value understood by the main function. This value can be either
        'auto' to detect cystein bridges automatically based on a default
        distance, 'none' to not have cystein bridges, or a float value to set
        cystein bridges based on a distance threshold.

    Raises
    ------
    argparse.ArgumentError
        Raised when the value cannot be converted.
    """
    try:
        result = float(value)
    except ValueError:
        lowered = value.lower()
        if lowered in ('auto', 'none'):
            return lowered
        raise argparse.ArgumentError(
            argument='-cys',
            message=('The value of the "cys" option must be "auto", "none", '
                     'or a distance in nanometers.'),
        )
    else:
        return result


def maxwarn(value):
    """
    Given a maxwarn specification, split it in a warning type, and the number
    to ignore.

    >>> maxwarn('3')
    (None, 3)
    >>> maxwarn('general:15')
    ('general', 15)
    >>> maxwarn('inconsistent-data')
    ('inconsistent-data, None)

    Parameters
    ----------
    value: str
        A warning type and a count, separated by a colon.

    Returns
    -------
    tuple[str, int]
        A warning type and the associated count to ignore. Either element can be
        None if not specified.

    Raises
    ------
    argparse.ArgumentTypeError
    """
    msg = ("Values for the -maxwarn option must be the name of a "
           "warning type, a number, or following the format "
           "'<warning-type>:<count>' where <warning-type> is the name "
           "of the warning type to ignore, and <count> is the number of "
           "warning of that type to ignore. "
           "'{value}' is not a valid value.".format(value=value))
    splitted = value.split(':')
    if len(splitted) == 1:
        try:
            count = int(value)
        except ValueError:
            # The value is not an int, so a warning type to ignore an
            # an unspecified number of
            return (value, None)
        else:
            return (None, count)
    elif len(splitted) == 2:
        try:
            count = int(splitted[1])
        except ValueError:
            pass  # The exception will be raised at the end of the function
        else:
            return (splitted[0], count)
    raise argparse.ArgumentTypeError(msg)


def entry():
    """
    Parses commandline arguments and performs the logic.
    """
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument('-V', '--version', action='version', version=VERSION)

    file_group = parser.add_argument_group('Input and output files')
    file_group.add_argument('-f', dest='inpath', required=True, type=Path,
                            help='Input file (PDB|GRO)')
    file_group.add_argument('-x', dest='outpath', required=True, type=Path,
                            help='Output coarse grained structure (PDB)')
    file_group.add_argument('-o', dest='top_path', type=Path,
                            help='Output topology (TOP)')
    file_group.add_argument('-sep', dest='keep_duplicate_itp',
                            action='store_true', default=False,
                            help='Write separate topologies for identical chains')
    file_group.add_argument('-merge', dest='merge_chains',
                            type=lambda x: x.split(','), action='append',
                            help='Merge chains: e.g. -merge A,B,C (+)')
    file_group.add_argument('-ignore', dest='ignore_res', nargs='+',
                            default=[], action='append', type=lambda x: x.split(','),
                            help='Ignore residues with that name: e.g. '
                            '-ignore HOH,LIG (+)')
    file_group.add_argument('-ignh', dest='ignore_h', default=False,
                            action='store_true',
                            help='Ignore all Hydrogen atoms in the input file')
    file_group.add_argument('-model', dest='modelidx', default=None, type=int,
                            help='Which MODEL to select. Only meaningful for'
                            ' PDB files.')
    file_group.add_argument('-bonds-from', dest='bonds_from',
                            choices=['name', 'distance', 'none', 'both'],
                            default='both',
                            help="How to determine connectivity in the input. "
                            "If 'none', only bonds from the input file (CONECT)"
                            " will be used.")
    file_group.add_argument('-bonds-fudge', dest='bonds_fudge', type=float,
                            default=1.2, help='Factor with which Van der Waals'
                            ' radii should be scaled when determining bonds '
                            'based on distances.')

    ff_group = parser.add_argument_group('Force field selection')
    ff_group.add_argument('-ff', dest='to_ff', default='martini3001',
                          help='Which forcefield to use')
    ff_group.add_argument('-from', dest='from_ff', default='universal',
                          help='Force field of the original structure.')
    ff_group.add_argument('-ff-dir', dest='extra_ff_dir', action='append',
                          type=Path, default=[],
                          help='Additional repository for custom force fields.')
    ff_group.add_argument('-map-dir', dest='extra_map_dir', action='append',
                          type=Path, default=[],
                          help='Additional repository for mapping files.')
    ff_group.add_argument('-list-ff', action='store_true', dest='list_ff',
                          help='List all known force fields, and exit.')
    ff_group.add_argument('-list-blocks', action='store_true', dest='list_blocks',
                          help='List all Blocks and Modifications known to the'
                          ' force field, and exit.')

    posres_group = parser.add_argument_group('Position restraints')
    posres_group.add_argument('-p', dest='posres', type=str.lower,
                              choices=('none', 'all', 'backbone'), default='none',
                              help='Output position restraints (none/all/backbone)')
    posres_group.add_argument('-pf', dest='posres_fc', type=float, default=1000,
                              help='Position restraints force constant in kJ/mol/nm^2')
    secstruct_group = parser.add_argument_group('Secondary structure handling')
    secstruct_exclusion = secstruct_group.add_mutually_exclusive_group()
    secstruct_exclusion.add_argument('-dssp', nargs='?', const='dssp',
                                     help='DSSP executable for determining structure')
    secstruct_exclusion.add_argument('-ss', dest='ss', type=str.upper,
                                     metavar='SEQUENCE',
                                     help=('Manually set the secondary '
                                           'structure of the proteins.'))
    secstruct_exclusion.add_argument('-collagen', action='store_true', default=False,
                                     help='Use collagen parameters')
    secstruct_group.add_argument('-ed', dest='extdih', action='store_true', default=False,
                                 help=('Use dihedrals for extended regions '
                                       'rather than elastic bonds'))

    rb_group = parser.add_argument_group('Protein elastic network')
    rb_group.add_argument('-elastic', action='store_true', default=False,
                          help='Write elastic bonds')
    rb_group.add_argument('-ef', dest='rb_force_constant', type=float, default=500,
                          help='Elastic bond force constant Fc in kJ/mol/nm^2')
    rb_group.add_argument('-el', dest='rb_lower_bound', type=float, default=0,
                          help='Elastic bond lower cutoff: F = Fc if rij < lo')
    rb_group.add_argument('-eu', dest='rb_upper_bound', type=float, default=0.9,
                          help='Elastic bond upper cutoff: F = 0  if rij > up')
    rb_group.add_argument('-ermd', dest='res_min_dist', type=int,
                          help=('The minimum separation between two residues to have an RB '
                                'the default value is set by the force-field.'),
                          default=None)
    rb_group.add_argument('-ea', dest='rb_decay_factor', type=float, default=0,
                          help='Elastic bond decay factor a')
    rb_group.add_argument('-ep', dest='rb_decay_power', type=float, default=1,
                          help='Elastic bond decay power p')
    rb_group.add_argument('-em', dest='rb_minimum_force', type=float, default=0,
                          help='Remove elastic bonds with force constant lower than this')
    rb_group.add_argument('-eb', dest='rb_selection',
                          type=lambda x: x.split(','), default=None,
                          help='Comma separated list of bead names for elastic bonds')
    rb_group.add_argument('-eunit', dest='rb_unit', default='molecule',
                          choices=['chain', 'molecule', 'all'],
                          help=('Establish what is the structural unit for the '
                                'elastic network. Bonds are only created within'
                                ' a unit.'))

    go_group = parser.add_argument_group('Virtual site based GoMartini')
    go_group.add_argument('-govs-includes', action='store_true', default=False,
                          help='Write include statements to use Vitrual Site Go Martini.')
    go_group.add_argument('-govs-moltype', default='molecule_0',
                          help=('Set the name of the molecule when using '
                                'Virtual Sites GoMartini.'))

    prot_group = parser.add_argument_group('Protein description')
    prot_group.add_argument('-scfix', dest='scfix',
                            action='store_true', default=False,
                            help='Apply side chain corrections.')
    prot_group.add_argument('-cys', dest='cystein_bridge',
                            type=_cys_argument,
                            default='none', help='Cystein bonds')
    prot_group.add_argument('-mutate', dest='mutations', action='append',
                            type=lambda s: s.split(':'), default=[],
                            help='Mutate a residue. Desired mutation is '
                            'specified as, e.g. A-PHE45:ALA. The format is '
                            '<chain>-<resname><resid>:<new resname>. Elements '
                            'of the specification can be omitted as required.')
    prot_group.add_argument('-modify', dest='modifications', action='append',
                            type=lambda s: s.split(':'), default=[],
                            help='Add a modification to a residue. Desired '
                            'modification is specified as, e.g. A-ASP45:ASP0. '
                            'The format is <chain>-<resname><resid>:<modification>.'
                            ' Elements of the specification can be omitted as '
                            'required.')
    prot_group.add_argument('-nter', dest='modifications', action='append',
                            type=lambda s: ['nter', s], default=[],
                            help='Shorthand for patching N-termini. An '
                            'N-terminus is defined as a residue which is '
                            'connected to 1 other residue, and has the highest '
                            'resid.')
    prot_group.add_argument('-cter', dest='modifications', action='append',
                            type=lambda s: ['cter', s], default=[],
                            help='Shorthand for patching C-termini. A '
                            'C-terminus is defined as a residue which is '
                            'connected to 1 other residue, and has the lowest '
                            'resid.')
    # Unfortunately there's no action=extend_const. append_const *almost* does
    # what we need, but it makes the resulting list too deep:
    # [[['cter', 'COOH-ter'], ['nter', 'NH2-ter']], ['ASP3', 'ASP0']]
    prot_group.add_argument('-nt', dest='neutral_termini',
                            action='store_true', default=False,
                            help='Set neutral termini (charged is default). '
                            'Alias for "-nter NH2-ter -cter COOH-ter"')

    debug_group = parser.add_argument_group('Debugging options')
    debug_group.add_argument('-write-graph', type=Path, default=None,
                             help='Write the graph as PDB after the MakeBonds step.')
    debug_group.add_argument('-write-repair', type=Path, default=None,
                             help=('Write the graph as PDB after the '
                                   'RepairGraph step. The resulting file may '
                                   'contain "nan" coordinates making it '
                                   'unreadable by most softwares.'))
    debug_group.add_argument('-write-canon', type=Path, default=None,
                             help=('Write the graph as PDB after the '
                                   'CanonicalizeModifications step. The '
                                   'resulting file may contain "nan" '
                                   'coordinates making it unreadable by most '
                                   'software.'))
    debug_group.add_argument('-v', dest='verbosity', action='count',
                             help='Enable debug logging output. Can be given '
                                  'multiple times.', default=0)
    debug_group.add_argument('-maxwarn', dest='maxwarn', type=maxwarn,
                             action='append', nargs='+', default=[],
                             help='The maximum number of allowed warnings. If '
                             'more warnings are encountered no output files are'
                             ' written.')

    args = parser.parse_args()
    if args.elastic and args.govs_includes:
        parser.error('A rubber band elastic network and GoMartini are not '
                     'compatible. The -elastic and -govs-include flags cannot '
                     'be used together.')

    if args.to_ff.startswith('elnedyn'):
        # FIXME: This type of thing should be added to the FF itself.
        LOGGER.info('The forcefield {} must always be used with an elastic '
                    'network. Enabling it now.', args.to_ff)
        args.elastic = True

    file_extension = args.inpath.suffix.upper()[1:]  # We do not keep the dot
    if file_extension in ['GRO'] and args.modelidx is not None:
        parser.error("GRO files don't know the concept of models.")
    if args.modelidx is None:
        # Set a sane default value. Can't do this using argparse machinery,
        # since we need to be able to check whether the flag was given.
        args.modelidx = 1

    bonds_from_name = args.bonds_from in ('name', 'both')
    bonds_from_dist = args.bonds_from in ('distance', 'both')

    loglevels = {0: logging.INFO, 1: logging.DEBUG, 2: 5}
    LOGGER.setLevel(loglevels[args.verbosity])

    known_force_fields = vermouth.forcefield.find_force_fields(
        Path(DATA_PATH) / 'force_fields'
    )
    known_mappings = read_mapping_directory(Path(DATA_PATH) / 'mappings',
                                            known_force_fields)

    # Add user force fields and mappings
    for directory in args.extra_ff_dir:
        try:
            vermouth.forcefield.find_force_fields(directory, known_force_fields)
        except FileNotFoundError:
            msg = '"{}" given to the -ff-dir option should be a directory.'
            raise ValueError(msg.format(directory))
    for directory in args.extra_map_dir:
        try:
            partial_mapping = read_mapping_directory(directory,
                                                     known_force_fields)
        except NotADirectoryError:
            msg = '"{}" given to the -map-dir option should be a directory.'
            raise ValueError(msg.format(directory))
        combine_mappings(known_mappings, partial_mapping)

    if args.list_ff:
        print('The following force fields are known:')
        for idx, ff_name in enumerate(reversed(list(known_force_fields)), 1):
            print('{:3d}. {}'.format(idx, ff_name))
        parser.exit()

    # Build self mappings
    partial_mapping = generate_all_self_mappings(known_force_fields.values())
    combine_mappings(known_mappings, partial_mapping)

    from_ff = args.from_ff
    if args.to_ff not in known_force_fields:
        raise ValueError('Unknown force field "{}".'.format(args.to_ff))
    if args.from_ff not in known_force_fields:
        raise ValueError('Unknown force field "{}".'.format(args.from_ff))
    #if from_ff not in known_mappings or args.to_ff not in known_mappings[from_ff]:
    #    raise ValueError('No mapping known to go from "{}" to "{}".'
    #                     .format(from_ff, args.to_ff))

    if args.list_blocks:
        print('The following Blocks are known to force field {}:'.format(args.from_ff))
        print(', '.join(known_force_fields[args.from_ff].blocks))
        print('The following Modifications are known to force field {}:'.format(args.from_ff))
        print(', '.join(known_force_fields[args.from_ff].modifications))
        print()
        print('The following Blocks are known to force field {}:'.format(args.to_ff))
        print(', '.join(known_force_fields[args.to_ff].blocks))
        print('The following Modifications are known to force field {}:'.format(args.to_ff))
        print(', '.join(known_force_fields[args.to_ff].modifications))
        parser.exit()


    # args.ignore_res is a pretty deep list: given "-ignore HOH CU,LIG -ignore LIG2"
    # it'll contain [[['HOH'], ['CU', 'LIG']], [['LIG2']]]
    ignore_res = set()
    for grp in args.ignore_res:
        ignore_res.update(*grp)

    if args.neutral_termini:
        args.modifications.append(['cter', 'COOH-ter'])
        args.modifications.append(['nter', 'NH2-ter'])
    else:
        if args.modifications:
            resspecs, mods = zip(*args.modifications)
        else:
            resspecs, mods = [], []
        if not any('cter' in resspec for resspec in resspecs):
            args.modifications.append(['cter', 'C-ter'])
        if not any('nter' in resspec for resspec in resspecs):
            args.modifications.append(['nter', 'N-ter'])

    # Reading the input structure.
    # So far, we assume we only go from atomistic to martini. We want the
    # input structure to be a clean universal system.
    # For now at least, we silently delete molecules with unknown blocks.
    system = read_system(args.inpath, ignore_resnames=ignore_res,
                         ignh=args.ignore_h, modelidx=args.modelidx)
    system = pdb_to_universal(
        system,
        delete_unknown=True,
        force_field=known_force_fields[from_ff],
        bonds_from_name=bonds_from_name,
        bonds_from_dist=bonds_from_dist,
        bonds_fudge=args.bonds_fudge,
        modifications=args.modifications,
        mutations=args.mutations,
        write_graph=args.write_graph,
        write_repair=args.write_repair,
        write_canon=args.write_canon,
    )

    LOGGER.info('Read input.', type='step')
    for molecule in system.molecules:
        LOGGER.debug("Read molecule {}.", molecule, type='step')

    target_ff = known_force_fields[args.to_ff]
    if args.dssp is not None:
        AnnotateDSSP(executable=args.dssp, savedir='.').run_system(system)
        AnnotateMartiniSecondaryStructures().run_system(system)
    elif args.ss is not None:
        AnnotateResidues(attribute='secstruct', sequence=args.ss,
                         molecule_selector=selectors.is_protein).run_system(system)
        AnnotateMartiniSecondaryStructures().run_system(system)
    elif args.collagen:
        if not target_ff.has_feature('collagen'):
            LOGGER.warning('The force field "{}" does not have specific '
                           'parameters for collagen (-collagen).',
                           target_ff.name, type='missing-feature')
        AnnotateResidues(attribute='cgsecstruct', sequence='F',
                         molecule_selector=selectors.is_protein).run_system(system)
    if args.extdih and not target_ff.has_feature('extdih'):
        LOGGER.warning('The force field "{}" does not define dihedral '
                       'angles for extended regions of proteins (-extdih).',
                       target_ff.name, type='missing-feature')
    vermouth.SetMoleculeMeta(extdih=args.extdih).run_system(system)
    if args.scfix and not target_ff.has_feature('scfix'):
        LOGGER.warning('The force field "{}" does not define angle and '
                       'torsion for the side chain corrections (-scfix).',
                       target_ff.name, type='missing-feature')
    vermouth.SetMoleculeMeta(scfix=args.scfix).run_system(system)

    ss_sequence = list(itertools.chain(*(
        dssp.sequence_from_residues(molecule, 'secstruct')
        for molecule in system.molecules
        if selectors.is_protein(molecule)
    )))

    if args.cystein_bridge == 'none':
        vermouth.RemoveCysteinBridgeEdges().run_system(system)
    elif args.cystein_bridge != 'auto':
        vermouth.AddCysteinBridgesThreshold(args.cystein_bridge).run_system(system)

    # Run martinize on the system.
    system = martinize(
        system,
        mappings=known_mappings,
        to_ff=known_force_fields[args.to_ff],
        delete_unknown=True,
    )

    # Apply position restraints if required.
    if args.posres != 'none':
        LOGGER.info('Applying position restraints.', type='step')
        node_selectors = {'all': selectors.select_all,
                          'backbone': selectors.select_backbone}
        node_selector = node_selectors[args.posres]
        vermouth.ApplyPosres(node_selector, args.posres_fc).run_system(system)

    if args.govs_includes:
        # The way Virtual Site GoMartini works has to be in sync with
        # Sebastian's create_goVirt.py script, until the method is fully
        # implemented in vermouth. One call of martinize2 must create a single
        # molecule, regardless of the number of fragments in the input.
        # The molecule type name is provided as an input with the -govs-moltype
        # flag to be consistent with the name provided to Sebastian's script.
        # The name cannot be guessed because a system may need to be composed
        # from multiple calls to martinize2 and create_goVirt.py.
        LOGGER.info('Adding includes for Virtual Site Go Martini.', type='step')
        LOGGER.info('The output topology will require files generated by '
                    '"create_goVirt.py".')
        vermouth.MergeAllMolecules().run_system(system)
        vermouth.SetMoleculeMeta(moltype=args.govs_moltype).run_system(system)
        vermouth.GoVirtIncludes().run_system(system)
        defines = ('GO_VIRT',)
    else:
        # Merge chains if required.
        if args.merge_chains:
            for chain_set in args.merge_chains:
                vermouth.MergeChains(chain_set).run_system(system)
        vermouth.NameMolType(deduplicate=not args.keep_duplicate_itp).run_system(system)
        defines = ()

    # Apply a rubber band elastic network is required.
    if args.elastic:
        LOGGER.info('Setting the rubber bands.', type='step')
        if args.rb_unit == 'molecule':
            domain_criterion = vermouth.processors.apply_rubber_band.always_true
        elif args.rb_unit == 'all':
            vermouth.MergeAllMolecules().run_system(system)
            domain_criterion = vermouth.processors.apply_rubber_band.always_true
        elif args.rb_unit == 'chain':
            domain_criterion = vermouth.processors.apply_rubber_band.same_chain
        else:
            message = 'Unknown value for -eunit: "{}".'.format(args.rb_unit)
            LOGGER.critical(message)
            raise ValueError(message)
        if args.rb_selection is not None:
            selector = functools.partial(
                selectors.proto_select_attribute_in,
                attribute='atomname',
                values=args.rb_selection,
            )
        else:
            selector = selectors.select_backbone
        rubber_band_processor = vermouth.ApplyRubberBand(
            lower_bound=args.rb_lower_bound,
            upper_bound=args.rb_upper_bound,
            decay_factor=args.rb_decay_factor,
            decay_power=args.rb_decay_power,
            base_constant=args.rb_force_constant,
            minimum_force=args.rb_minimum_force,
            selector=selector,
            domain_criterion=domain_criterion,
            res_min_dist=args.res_min_dist
        )
        rubber_band_processor.run_system(system)

    LOGGER.info('Writing output.', type='step')
    for molecule in system.molecules:
        LOGGER.debug("Writing molecule {}.", molecule, type='step')

    # Write the topology if requested
    # grompp has a limit in the number of character it can read per line
    # (due to the size limit of a buffer somewhere in its implementation).
    # The command line can be longer than this limit and therefore
    # prevent grompp from reading the topology.
    gromacs_char_limit = 4000  # the limit is actually 4095, but I play safe
    command = ' '.join(sys.argv)
    if len(command) > gromacs_char_limit:
        command = command[:gromacs_char_limit] + ' ...'
    header = [
        'This file was generated using the following command:',
        command,
        VERSION,
    ]
    if None not in ss_sequence:
        header += [
            'The following sequence of secondary structure ',
            'was used for the full system:',
            ''.join(ss_sequence),
        ]

    if args.top_path is not None:
        write_gmx_topology(system, args.top_path, defines=defines, header=header)

    # Write a PDB file.
    vermouth.pdb.write_pdb(system, str(args.outpath), omit_charges=True)

    # TODO: allow ignoring warnings per class/amount (i.e. ignore 2
    #       inconsistent-data warnings)

    leftover_warnings = ignore_warnings_and_count(COUNTER, args.maxwarn)
    if leftover_warnings:
        LOGGER.error('{} warnings were encountered after accounting for the '
                     '-maxwarn flag. No output files will be '
                     'written. Consider fixing the warnings, or if you are sure'
                     ' they are harmless, use the -maxwarn flag.', leftover_warnings)
        sys.exit(2)
    else:
        DeferredFileWriter().write()
        vermouth.Quoter().run_system(system)


if __name__ == '__main__':
    entry()
