"""
Gromacs support in OpenPathSampling

OpenPathSampling supports Gromacs as an "external engine," meaning that
Gromacs itself runs as a coprocess, and OPS uses the file system as an
intermediary to monitor when a path needs to be terminated.
"""

import logging
logger = logging.getLogger(__name__)

from mdtraj.formats import TRRTrajectoryFile
import mdtraj as md

from openpathsampling.engines import ExternalEngine
from openpathsampling.engines import features
from openpathsampling.engines.snapshot import BaseSnapshot, SnapshotDescriptor
from openpathsampling.engines.topology import MDTrajTopology
from openpathsampling.engines.external_snapshots import \
        ExternalMDSnapshot, InternalizedMDSnapshot
from openpathsampling.tools import ensure_file

import os
import psutil
import shlex
import time
import numpy as np

from openpathsampling.engines.external_engine import (
    _debug_open_files, close_file_descriptors
)

def _remove_file_if_exists(filename):  # pragma: no cover
    #  not requiring coverage here because it's part of Gromacs integration;
    #  gets covered if gmx is installed though
    if os.path.isfile(filename):
        os.remove(filename)

class _GroFileEngine(ExternalEngine):
    SnapshotClass = ExternalMDSnapshot
    InternalizedSnapshotClass = InternalizedMDSnapshot
    def __init__(self, gro):
        self.gro = gro
        traj = md.load(gro)
        self.topology = MDTrajTopology(traj.topology)
        n_atoms = self.topology.n_atoms
        n_spatial = self.topology.n_spatial
        descriptor = SnapshotDescriptor.construct(
            snapshot_class=self.SnapshotClass,
            snapshot_dimensions={'n_spatial': n_spatial,
                                 'n_atoms': n_atoms}
        )
        super(_GroFileEngine, self).__init__(options={},
                                            descriptor=descriptor,
                                            template=None)

    def to_dict(self):
        return {'gro': self.gro}

    @classmethod
    def from_dict(cls, dct):
        return cls(dct['gro'])

    def read_frame_data(self, file_name, file_position):
        traj = md.load(file_name)
        xyz = traj.xyz[0]
        vel = np.zeros(shape=xyz.shape)
        box = traj.unitcell_vectors[0]
        return (xyz, vel, box)



def snapshot_from_gro(gro_file):
    template_engine = _GroFileEngine(gro_file)
    snapshot = ExternalMDSnapshot(file_name=gro_file,
                                  file_position=0,
                                  engine=template_engine)
    return snapshot


class GromacsEngine(ExternalEngine):
    """
    External engine wrapper for Gromacs (using indirect API).

    This provides Gromacs support, using our indirect engine API (TODO
    link).

    OPS runs Gromacs based on the mdp file that you provide. This mdp file
    MUST output in the TRR format, and the velocity and position save
    frequency in the TRR must be the same (that is, you need to have
    ``nstxout`` = ``nstvout``, and they must not be 0).

    Parameters
    ----------
    gro : string
        File for the grompp ``-c`` flag. This is often a .gro, but note that
        you may get better support with other integrations (e.g., MDTraj) if
        you use a PDB.
    mdp : string
        .mdp file
    top : string
        .top file
    options : dict
        Dictionary of option name to value. Gromacs-specific option names
        are

            * ``gmx_executable``: Prefix to gromacs commands, which are run
              as ``{gmx_prefix}command``. Default is 'gmx ' (note the
              space).  This allows you to use either Gromacs 4 or Gromacs 5,
              as well as specifying the path to your version of Gromacs.
            * ``grompp_args``: Additional arguments to ``grompp``. The
              defaults take ``-c {self.gro} -f {self.mdp} -p {self.top}
              -t {self.input_file}``, where the input filename is set by
              :meth:`.set_filenames`. Default is the empty string.
            * ``mdrun_args``: Additional arguments to ``mdrun``. The
              defaults take ``-s topol.top -o self.output_file
              -e self.edr_file -g self.log_file``, where the ``topol.top``
              is generated by :meth:`.prepare`, and the other filenames are
              set by :meth:`.set_filenames`. Default is the empty string.
            * ``snapshot_timestep``: time between frames analysed by ops.
              You keep track of the unit, I'd advise ps so the output
              rates will be in ps. Example. 2 fs timestep in the mdp with
              nstxout of 30 would give snapshot_timestep of 60 fs = 0.06 ps

    base_dir : string
        root directory where all files will be found (defaults to pwd)
    prefix : string
        prefix within ``base_dir`` for output folders (defaults to gmx)
    """
    _default_options = dict(ExternalEngine._default_options,
        **{
            'gmx_executable': "gmx ",
            'grompp_args': "",
            'mdrun_args': "",
            'snapshot_timestep':1.0
        }
    )
    GROMPP_CMD = ("{e.options[gmx_executable]}grompp -c {e.gro} "
                  + "-f {e.mdp} -p {e.top} -t {e.input_file} "
                  + "-po {e.mdout_file} -o {e.tpr_file} "
                  + "{e.options[grompp_args]}")
    MDRUN_CMD = ("{e.options[gmx_executable]}mdrun -s {e.tpr_file} "
                 + "-o {e.output_file} -e {e.edr_file} -g {e.log_file} "
                 + "{mdrun_args}")
    # use these as CMD.format(e=engine, **engine.options)
    SnapshotClass = ExternalMDSnapshot
    InternalizedSnapshotClass = InternalizedMDSnapshot
    def __init__(self, gro, mdp, top, options, base_dir="", prefix="gmx"):
        self.base_dir = base_dir
        self.gro = os.path.join(base_dir, gro)
        self.mdp = os.path.join(base_dir, mdp)
        self.top = os.path.join(base_dir, top)
        self.prefix = os.path.join(base_dir, prefix)

        self.gro_contents, self._gro_hash = ensure_file(self.gro)
        self.mdp_contents, self._mdp_hash = ensure_file(self.mdp)
        self.top_contents, self._top_hash = ensure_file(self.top)

        # TODO: move to a later stage, before first traj
        dirs = [self.prefix + s for s in ['_trr', '_log', '_edr']]
        for d in dirs:
            try:
                os.mkdir(d)
            except OSError:
                pass  # the directory already exists

        self._file = None  # file open/close efficiency
        self._last_filename = None
        # TODO: add snapshot_timestep; first via options, later read mdp
        template = snapshot_from_gro(self.gro)
        self.topology = template.topology
        descriptor = template.engine.descriptor  # descriptor from gro file

        # initial placeholders
        self.input_file = "INITIAL.trr"
        self.output_file = self.prefix + "_trr/OUTPUT_NAME.trr"
        self.edr_file = self.prefix + "_edr/OUTPUT_NAME.edr"
        self.log_file = self.prefix + "_log/OUTPUT_NAME.log"
        self.tpr_file = "topol.tpr"
        self.mdout_file = "mdout.mdp"

        self._mdtraj_topology = None

        super(GromacsEngine, self).__init__(options, descriptor, template,
                                             first_frame_in_file=True)

    def to_dict(self):
        dct = super(GromacsEngine, self).to_dict()
        local_dct = {
            'gro': self.gro,
            'mdp': self.mdp,
            'top': self.top,
            'options': self.options,
            'base_dir': self.base_dir,
            'prefix': self.prefix,
            'gro_contents': self.gro_contents,
            'mdp_contents': self.mdp_contents,
            'top_contents': self.top_contents,
        }
        dct.update(local_dct)
        return dct

    @classmethod
    def from_dict(cls, dct):
        dct = dict(dct)  # make a copy
        for ftype in ['gro', 'top', 'mdp']:
            contents = dct.pop(ftype + "_contents")
            _ = ensure_file(filename=dct[ftype], old_contents=contents)
        return super(GromacsEngine, cls).from_dict(dct)

    @property
    def mdtraj_topology(self):
        if self._mdtraj_topology:
            return self._mdtraj_topology
        return self.topology.mdtraj

    @mdtraj_topology.setter
    def mdtraj_topology(self, value):
        self._mdtraj_topology = value

    def read_frame_data(self, filename, frame_num):
        """
        Returns pos, vel, box or raises error
        """
        # if self._last_filename != filename:
            # try:
                # self._file.close()
            # except AttributeError:
                # pass  # first time thru, self._file is None
            # self._file = TRRTrajectoryFile(filename)
        # f = self._file
        # do we need to reopen the TRR each time to avoid problems with the
        # fiel length changing?
        trr = TRRTrajectoryFile(filename)
        f = trr
        logger.debug("Reading file %s frame %d (of %d)",
                     filename, frame_num, len(f))
        # logger.debug("File length: %d", len(f))
        try:
            f.seek(offset=frame_num)
            data = f._read(n_frames=1, atom_indices=None, get_velocities=True)
        finally:
            trr.close()

        return data[0][0], data[5][0], data[3][0]

    def read_frame_from_file(self, file_name, frame_num):
        # note: this only needs to return the file pointers -- but should
        # only do so once that frame has been written!
        basename = os.path.basename(file_name)
        # basename should be in the format [0-9]+\.trr (as set by the
        # trajectory_filename method)
        # file_number = int(basename.split('.')[0])
        try:
            xyz, vel, box = self.read_frame_data(file_name, frame_num)
        except (IndexError, OSError, IOError) as e:
            # this means that no such frame exists yet, so we return None
            # IndexError in older version, OSError more recently (specific
            # MDTraj error)
            logger.debug("Expected exception caught: " + str(e))
            close_file_descriptors(basename)
            return None
        except RuntimeError as e:
            # TODO: matches "TRR read error"
            logger.debug("Received partial frame for %s %d", file_name,
                         frame_num+1)
            return 'partial'
        else:
            logger.debug("Creating snapshot")
            snapshot =  ExternalMDSnapshot(file_name=file_name,
                                           file_position=frame_num,
                                           engine=self)
            return snapshot

    def write_frame_to_file(self, filename, snapshot, mode='w'):
        if os.path.isfile(filename):
            # stop if we already have this file; could also happen because
            # of a weird behavior in a mover. You must remove the files if
            # you don't want them.
            raise RuntimeError("File " + str(filename) + " exists. "
                               + "Preventing overwrite.")
        # type control before passing things to Cython code
        xyz = np.asarray([snapshot.xyz], dtype=np.float32)
        time = np.asarray([0.0], dtype=np.float32)
        step = np.asarray([0], dtype=np.int32)
        box = np.asarray([np.asarray(snapshot.box_vectors)], dtype=np.float32)
        lambd = np.asarray([0.0], dtype=np.float32)
        vel = np.asarray([np.asarray(snapshot.velocities)], dtype=np.float32)

        try:
            trr = TRRTrajectoryFile(filename, mode)
            trr._write(xyz, time, step, box, lambd, vel)
        finally:
            trr.close()
            close_file_descriptors(filename)

    def trajectory_filename(self, number):
        # TODO: remove code path allowing ints (API break for 2.0)
        trr_dir = self.prefix + "_trr/"
        if isinstance(number, int):
            num_str = num_str = '{:07d}'.format(number)
        else:
            num_str = number
        return trr_dir + num_str + '.trr'

    def set_filenames(self, number):
        if isinstance(number, int):
            num_str = '{:07d}'.format(number + 1)
            self.output_file = self.trajectory_filename(number + 1)
            init_filename = "initial_frame.trr"
            self.filename_setter.reset(number + 1)
        else:
            num_str = number
            self.output_file = self.trajectory_filename(num_str)
            init_filename = num_str + "_initial_frame.trr"
            self.mdout = num_str + "_mdout.mdp"
            self.tpr_file = num_str + "_" + "topol.tpr"

        self.input_file = os.path.join(self.base_dir, init_filename)
        self.edr_file = os.path.join(self.prefix + "_edr", num_str + '.edr')
        self.log_file = os.path.join(self.prefix + "_log", num_str + '.log')

    @property
    def grompp_command(self):
        cmd = self.GROMPP_CMD.format(e=self)
        return cmd

    def prepare(self):  # pragma: no cover
        # coverage ignored b/c Travis won't have gmx. However, we do have a
        # test that covers this if gmx is present (otherwise it is skipped)

        # ensure that the files haven't changed before we run trajectory
        _ = ensure_file(self.gro, self.gro_contents, self._gro_hash)
        _ = ensure_file(self.mdp, self.mdp_contents, self._mdp_hash)
        _ = ensure_file(self.top, self.top_contents, self._top_hash)

        # grompp and mdrun
        cmd = self.grompp_command
        logger.info(cmd)
        run_cmd = shlex.split(cmd)
        return_code = psutil.Popen(run_cmd, preexec_fn=os.setsid).wait()
        return return_code

    def cleanup(self):  # pragma: no cover
        # tested when traj is run, which we don't on CI
        _remove_file_if_exists(self.input_file)
        _remove_file_if_exists(self.tpr_file)
        _remove_file_if_exists(self.mdout_file)

    def engine_command(self):
        args = self.options['mdrun_args'].format(prev_traj=self._traj_num-1,
                                                 next_traj=self._traj_num)
        cmd = self.MDRUN_CMD.format(e=self, mdrun_args=args)
        return cmd
