#!/usr/bin/env python3
"""
Module for working with NSG. See https://github.com/OpenSourceBrain/pynsgr

File: pyneuroml/nsgr.py

Copyright 2024 NeuroML contributors
"""

import logging
import os
import pathlib
import typing
from zipfile import ZipFile

from pyneuroml.runners import generate_sim_scripts_in_folder

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

try:
    from pynsgr.commands.nsgr_submit import nsgr_submit
except ImportError:
    logger.warning("Please install optional dependencies to use NSG features:")
    logger.warning("pip install pyneuroml[nsgr]")


def run_on_nsg(
    engine: str,
    lems_file_name: str,
    root_dir: typing.Optional[str] = None,
    nsg_sim_config: typing.Dict[typing.Any, typing.Any] = {},
    run_dir: typing.Optional[str] = None,
    dry_run: bool = False,
    *engine_args: typing.Any,
    **engine_kwargs: typing.Any,
):
    """Run NeuroML/LEMS simulation on NSG.

    Since the versions of tools on NSG may differ from what you have installed
    locally, this method first generates the simulation engine specific files
    (runner script for NEURON and mod files, for example) for the provided
    engine in a new folder, zips them up, writes the config files, and then
    submits the job to NSG using pynsgr.

    The data generated by the job (if the command is not interrupted) is also
    downloaded to the newly created directory.

    Please ensure that you have set up an account and have your NSG
    configuration file populated as noted in pynsgr.

    - https://nsgr.sdsc.edu:8443/restusers/documentation
    - https://nsgr.sdsc.edu:8443/restusers/docs/tools

    Default for the nsg_sim_config is below, keys provided by the user in
    nsg_sim_config overwrite these:

    .. code:: python

        nsg_sim_config_dict = {
            "number_cores_": "1",
            "number_nodes_": "1",
            "number_gbmemorypernode_": "1",
            "tasks_per_node_": "1",
            "runtime_": "0.5",
            'toolId': "PY_EXPANSE",
            'nrnivmodl_o_': "1",
            "cmdlineopts_": "-nogui"
        }

    .. versionadded:: 1.0.10

    :param engine: name of engine: suffixes of the run_lems_with functions
    :type engine: str
    :param lems_file_name: name of LEMS simulation file
    :type lems_file_name: str
    :param root_dir: directory in which LEMS simulation file lives
        Any included files must be relative to this main directory
    :type root_dir: str
    :param nsg_sim_config: dict containing params and values that will be
        printed to testParam.properties
    :type nsg_sim_config: dict
    :param run_dir: directory in which model files are copied, backend specific
        files are generated, and from which the archive is sent to NSG. Results
        from the NSG run will also be downloaded to this directory.

        By default, this is the directory that the command is run from (".")

        It is good practice to separate directories where simulations are run
        from the source of the model/simulations.
    :type run_dir: str
    :param engine_args: positional args to be passed to the engine runner
        function
    :param engine_kwargs: keyword args to be be passed to the engine runner
        function
    :param dry_run: do everything but do not submit
    :type dry_run: bool
    :returns: name of new directory
    :rtype: str
    """
    supported_engines = ["jneuroml_neuron", "jneuroml_netpyne"]
    if engine not in supported_engines:
        print(f"Engine {engine} is not currently supported on NSG")
        print(f"Supported engines are: {supported_engines}")
        return

    logger.debug(f"NSGR: engine is {engine}")

    zipfile_name = lems_file_name.replace(".xml", "") + "_NSG.zip"
    # default dictionary
    nsg_sim_config_dict = {
        "number_cores_": "1",
        "number_nodes_": "1",
        "tasks_per_node_": "1",
        "number_gbmemorypernode_": "1",
        "runtime_": "0.5",
        "toolId": "PY_EXPANSE",
        "nrnivmodl_o_": "1",
        "cmdlineopts_": "-nogui",
    }

    # update dict based on user values
    for key, val in nsg_sim_config.items():
        nsg_sim_config_dict[key] = val

    # NSG requires that the top level directory exist
    nsg_dir = pathlib.Path(zipfile_name.replace(".zip", ""))

    cwd = pathlib.Path.cwd()

    tdir = generate_sim_scripts_in_folder(
        engine=engine,
        lems_file_name=lems_file_name,
        root_dir=root_dir,
        run_dir=run_dir,
        generated_files_dir_name=str(nsg_dir),
    )

    logger.info("Generating zip file")
    runner_file = ""

    os.chdir(str(tdir))
    generated_files = os.listdir(nsg_dir)

    print(f"Generated files are {generated_files}")

    with ZipFile(zipfile_name, "w") as archive:
        for f in generated_files:
            if engine == "jneuroml_neuron":
                if f.endswith("_nrn.py"):
                    runner_file = f
            elif engine == "jneuroml_netpyne":
                if f.endswith("_netpyne.py"):
                    runner_file = f
            fpath = pathlib.Path(f)
            moved_path = nsg_dir / fpath
            archive.write(str(moved_path))

    logger.debug("Printing testParam.properties")
    nsg_sim_config_dict["filename_"] = runner_file
    logger.debug(f"NSG sim config is: {nsg_sim_config_dict}")

    with open("testParam.properties", "w") as file:
        for key, val in nsg_sim_config_dict.items():
            print(f"{key}={val}", file=file)

    logger.debug("Printing testInput.properties")
    with open("testInput.properties", "w") as file:
        print(f"infile_=@./{zipfile_name}", file=file)

    print(f"{zipfile_name} generated")
    # uses argv, where the first argument is the script itself, so we must pass
    # something as the 0th index of the list
    if not dry_run:
        if nsgr_submit(["", ".", "validate"]) == 0:
            print("Attempting to submit to NSGR")
            return nsgr_submit(["", ".", "run"])
    else:
        print("Dry run mode enabled. Not submitting to NSG.")

    os.chdir(str(cwd))
    return tdir
