#!/usr/bin/env python
"""CLI utility to run gmx genion with AiiDA.

Usage: gmx_genion --help
"""

import os

import click

from aiida import cmdline, engine
from aiida.plugins import CalculationFactory, DataFactory

from aiida_gromacs import helpers
from aiida_gromacs.utils import searchprevious


def launch(params):
    """Run genion.

    Uses helpers to add gromacs on localhost to AiiDA on the fly.
    """

    # Prune unused CLI parameters from dict.
    params = {k:v for k,v in params.items() if v != None}

    # dict to hold our calculation data.
    inputs = {
        "metadata": {
            "description": params.pop("description"),
        },
    }

    # If code is not initialised, then setup.
    if "code" in inputs:
        inputs["code"] = params.pop("code")
    else:
        computer = helpers.get_computer()
        inputs["code"] = helpers.get_code(entry_point="gromacs", computer=computer)
        inputs["code"] = helpers.get_code(entry_point="bash", computer=computer)


    # Prepare input parameters in AiiDA formats.
    SinglefileData = DataFactory("core.singlefile")
    inputs["tprfile"] = SinglefileData(file=os.path.join(os.getcwd(), params.pop("s")))
    inputs["topfile"] = SinglefileData(file=os.path.join(os.getcwd(), params.pop("p")))

    if "n" in params:
        inputs["n_file"] = SinglefileData(file=os.path.join(os.getcwd(), params.pop("n")))

    GenionParameters = DataFactory("gromacs.genion")
    inputs["parameters"] = GenionParameters(params)

    # check if inputs are outputs from prev processes
    inputs = searchprevious.get_prev_inputs(inputs, ["tprfile", "topfile"])


    # check if a pytest test is running, if so run rather than submit aiida job
    # Note: in order to submit your calculation to the aiida daemon, do:
    # pylint: disable=unused-variable
    if "PYTEST_CURRENT_TEST" in os.environ:
        future = engine.run(CalculationFactory("gromacs.genion"), **inputs)
    else:
        future = engine.submit(CalculationFactory("gromacs.genion"), **inputs)


@click.command()
@cmdline.utils.decorators.with_dbenv()
@cmdline.params.options.CODE()
# Plugin options
@click.option("--description", default="record genion data provenance via the aiida_gromacs plugin", type=str, help="Short metadata description")
# Input file options
@click.option("-s", default="topol.tpr", type=str, help="Input structure file")
@click.option("-n", type=str, help="Index file")
@click.option("-p", default="topol.top", type=str, help="Topology file")
# Output file options
@click.option("-o", default="out.gro", type=str, help="Output structure file")
# Other parameter options
@click.option("-np", type=str, help="Number of positive ions")
@click.option("-pname", default="NA", type=str, help="Name of positive ion")
@click.option("-pq", type=str, help="Charge of the positive ion")
@click.option("-nn", type=str, help="Number of negative ions")
@click.option("-nname", default="CL", type=str, help="Name of negative ion")
@click.option("-nq", type=str, help="Charge of the negative ion")
@click.option("-rmin", type=str, help="Minimum distance between ions and non-solvent")
@click.option("-seed", type=str, help="Seed for random number generator (0 means generate)")
@click.option("-conc", type=str, help="Specify salt concentration (mol/liter). This will add sufficient ions to reach up to the specified concentration as computed from the volume of the cell in the input .tpr file. Overrides the -np and -nn options.")
@click.option(
    "-neutral", default="false", type=str, help="Neutralise the system with ions"
)
def cli(*args, **kwargs):
    # pylint: disable=unused-argument
    # pylint: disable=line-too-long
    """Run example.

    Example usage:

    $ gmx_genion --code gmx@localhost -s 1AKI_ions.tpr -p 1AKI_topology.top -pname NA -nname CL -neutral true -o 1AKI_solvated_ions.gro

    Alternative (automatically tried to create gmx@localhost code, but requires
    gromacs to be installed and available in your environment path):

    $ gmx_genion -s 1AKI_ions.tpr -p 1AKI_topology.top -pname NA -nname CL -neutral true -o 1AKI_solvated_ions.gro

    Help: $ gmx_genion --help
    """

    launch(kwargs)


if __name__ == "__main__":
    cli()  # pylint: disable=no-value-for-parameter
