r"""TaskInterface class corresponding to the program DL_MONTE"""


import logging
import os
import shutil

import dlmontepython.simtask.task as task

import dlmontepython.htk.sources.dlmonte as dlmonte
import dlmontepython.htk.sources.dlconfig as dlconfig
import dlmontepython.htk.sources.dlmove as dlmove




logger = logging.getLogger(__name__)




class DLMonteInterface(task.TaskInterface):

    r"""TaskInterface class corresponding to the program DL_MONTE

    TaskInterface class corresponding to the program DL_MONTE. This
    uses Kevin Stratford's DL_MONTE Python toolkit to do the heavy 
    lifting.

    Attributes
    ----------
    runner : htk.sources.dlmonte.DLMonteRunner
        htk.sources.dlmonte.DLMonteRunner object used to execute 
        DL_MONTE simulations. In the constructor for this class the
        path of the executable is used to initialise `runner`

    """

    def __init__(self, executable):

        r"""Constructor

        Constructor

        Parameters
        ----------
        executable : str
            The path to the DL_MONTE executable

        """

        self.runner = dlmonte.DLMonteRunner(executable)




    def copy_input_files(self, fromdir, todir):

        r"""Copy simulation input files between directories

        Copy simulation input files between directories

        Parameters
        ----------
        fromdir : str
            The directory containing the input files
        todir : str
            The target directory for the input files
       
        """

        siminput = dlmonte.DLMonteInput.from_directory(fromdir)
        siminput.to_directory(todir)




    def run_sim(self, simdir):

        r"""Run a simulation

        Run a simulation in a specified directory

        Parameters
        ----------
        simdir : str
            The directory in which the simulation is to be executed
       
        """

        self.runner.directory = simdir
        self.runner.execute()


    

    def resume_sim(self, oldsimdir, simdir):

        r"""Resume a simulation from a checkpointed state

        Resume a simulation whose checkpointed state is located in the
        directory `oldsimdir`, and run the resumed/new simulation in
        the directory `simdir`

        Parameters
        ----------
        oldsimdir : str
            The directory in which the files from the 'old' simulation
            to be resumed reside
        simdir : str
            The directory in which the new/resumed simulation is to be
            executed

        """

        # Copy the input files

        # Ideal code for copying input files, but the REVCON bug below stops it working...
        #siminput = dlmonte.DLMonteInput.from_directory(oldsimdir)
        #siminput.to_directory(simdir)

        # A simple workaround...
        shutil.copyfile(oldsimdir+"/CONTROL",simdir+"/CONTROL")
        shutil.copyfile(oldsimdir+"/FIELD",simdir+"/FIELD")

        # Copy the REVCON configuration to CONFIG

        ##################################################################################
        # This is the prefered code for using the REVCON configuration from the last
        # simulation as the CONFIG for the next simulation. However the dlconfig.from_file 
        # function cannot cope with the 'EXTRAS' on line 2 of REVCON
        #
        # If we are beyond the 1st simulation use the REVCON output file from the previous
        # simulation as the CONFIG for this simulation
        #if i>0:
        #    prevsimdir=self.simdir_header+str(i)
        #    siminput.config=dlconfig.from_file(prevsimdir+"/REVCON.000")
        ###################################################################################

        shutil.copyfile(oldsimdir+"/REVCON.000",simdir+"/CONFIG")

        self.run_sim(simdir)




    def extract_data(self, observable, simdir):

        r"""Extract simulation data corresponding to a specified observable
        
        Extract data corresponding to a specified observable from output
        files generated by a simulation in the directory `simdir`

        Parameters
        ----------
        observable : task.Observable
            An observable
        simdir : str
            The directory containing the output files to extract data from

        Returns
        -------
        array
            An array containing one or more values corresponding to 
            `observable`.
       
        Raises
        ------
        ValueError
            If `observable` is not recognised or supported by this function

        Notes
        -----
        * Normally the returned array would be a time series, e.g., if
          `observable` is "energy" then the returned array would contain
          the observed values of the energy vs. simulation time obtained
          from the simulation. However the array need not necessarily be
          a time series. E.g. it could be a radial distribution function.
          The nature of the values will depend on `observable`

        """

        # Load data - including YAMLDATA
        # After this 'simoutput.yamldata.metadata' contains the metadata
        # and 'simoutput.yamldata.data' contains the data
        simoutput = dlmonte.DLMonteOutput.load(simdir)
        
        # Gather data corresponding to the observable of interest
        data = []
        for frame in simoutput.yamldata.data:

            # How to map the elements of observable.descriptor to the
            # appropriate element in 'frame'...
            #
            # To get, e.g., frame['energy'], assume observable.descriptor = ('energy',),
            # and that len(observable.descriptor)==1
            #
            # To get, e.g., frame['nmol'][2], assume observable.descriptor = ('nmol',2),
            # and that len(observable.descriptor)==2.
            #

            if len(observable.descriptor) == 1:
                
                data.append(frame[observable.descriptor[0]])

            elif len(observable.descriptor) == 2:
                
                data.append( frame[observable.descriptor[0]][observable.descriptor[1]] )

            elif len(observable.descriptor) > 2:

                raise NotImplementedError("Depth in YAML frame greater than 2 not supported")

        return data




    def amend_input_parameter(self, dir, param, value):


        r"""Amend a specific simulation input parameter

        Amend the parameter `parameter` in the input files in `dir` to 
        take the value `value`

        Parameters
        ----------
        dir : str
            The directory containing the input files
        parameter : str
            Name of parameter to amend
        value : str
            New value of the parameter

        Raises
        ------
        ValueError
            If `parameter` is not recognised or supported by this function

        """

        # This is not general and a bit scrappy, and should be improved
        
        # As a special case catch the thermodynamic activity or chemical potential in DL_MONTE
        # for a molecule. THIS WORKS ON THE FIRST MOLECULE ONLY IN GCMC INSERT MOLECULE MOVES.
        if param=="molchempot":
            
            siminput = dlmonte.DLMonteInput.from_directory(dir)
           
            # The relevant location is siminput.control.main_block.moves.
            # Search through the moves and locate the molecule GCMC move
            for move in siminput.control.main_block.moves:
                
                if isinstance(move,dlmove.InsertMoleculeMove):
                    move.movers[0]["molpot"] = value
                    
            siminput.to_directory(dir)                  

        else:
        
            # Amend CONTROL file for now only
            siminput = dlmonte.DLMonteInput.from_directory(dir)
            siminput.control.main_block.statements[param] = value
            siminput.to_directory(dir)

        
