#!/usr/bin/python
# -*- coding: utf-8 -*-

#    Copyright 1998-2014 by Paweł T. Jochym <pawel.jochym@ifj.edu.pl>
#
#    qe-util is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.

#    qe-util is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.

#    You should have received a copy of the GNU General Public License
#    along with QE-util.  If not, see <http://www.gnu.org/licenses/>.


import os
import sys
from glob import glob
from os.path import join, isfile, islink
import subprocess
import time

import numpy as np

import ase.units
from ase.units import Bohr, Hartree, Rydberg, fs
from ase import Atom, Atoms
from ase.data import atomic_numbers, chemical_symbols
from ase.calculators.calculator import Calculator, FileIOCalculator, Parameters, kpts2mp, ReadError, all_changes

# Doubtful practice should be more selective in future
from readers import *
from writers import *
from qeio import *
from analyzers import *

import copy

# The exception for the calc runnin but not ready.
class CalcNotReadyError(Exception):
    pass


class QuantumEspresso(FileIOCalculator):
    """Class for doing Quantum Espresso calculations.

    The default parameters are very close to those that 
    the QE Fortran code would use.  These are the exceptions::
    
    The location of the pseudo potentials is by default
      /usr/share/espresso/pseudo
    or set from the ESPRESSO_PP_PATH environment variable
    
    The executables are by default pw.x, dos.x ph.x matdyn.x and q2r.x
    The calculation can use mpi parallelisation with property 'procs'.
    
    Supported parameters:
      'kpts','use_symmetry','procs','label'
      CONTROL
        'calc','pseudo_dir'
        'tstress', 'tprnfor','nstep','pseudo_dir','outdir',
        'wfcdir', 'prefix','forc_conv_thr', 'etot_conv_thr'
      SYSTEM
        'ecutwfc','ibrav','nbnd','occupations','degauss','smearing','ecutrho','nbnd'
      ELECTRONS
        'conv_thr','mixing_beta','mixing_mode','diagonalization',
        'mxing_ndim','electron_maxstep'
      

    """

    implemented_properties = ['energy', 'forces', 'stress','phdos','edos']
    
    pw_cmd='mpiexec -n %(procs)d  pw.x <%(iofile)s.in >%(iofile)s.out'
    edos_cmd='mpiexec -n %(procs)d  dos.x <edos_pp.in >edos_pp.out'
    ph_cmd='mpiexec -n %(procs)d  ph.x <ph.in >ph.out'
    matdyn_cmd='mpiexec -n %(procs)d matdyn.x <matdyn.in >matdyn.out'
    phdos_cmd= 'mpiexec -n %(procs)d matdyn.x <phdos.in  >phdos.out'
    q2r_cmd='q2r.x <q2r.in >q2r.out'

    cmds={
        'scf': pw_cmd,
        'ph': ph_cmd,
        'matdyn': matdyn_cmd,
        'phdos': phdos_cmd,
        'q2r': q2r_cmd,
        'edos': edos_cmd
    }

    default_parameters = {
                'calc':'scf',
                'outdir': 'tmp',
                'wfcdir': 'tmp',
                'asr':'crystal',
                'pseudo_dir': '/usr/share/espresso/pseudo',
                'ndos': 200,
                'points': 100,
                'tstress': True,
                'tprnfor': True,
                'ecutwfc': 50,
                'ibrav': 0,
                'pp_type': 'nc',
                'pp_format': 'UPF',
                'use_symmetry': True,
                'kpt_type': 'automatic',
                'kpt_shift': [0,0,0],
                'procs': 1
    }
    'Default parameters'

    def __init__(self, label=None, atoms=None, wdir='./', restart=None, **kwargs):
        FileIOCalculator.__init__(self,label=label,atoms=atoms,command=None,**kwargs)
        self.label=label
        self.prefix=label
        self.wdir=wdir
        if restart :
            # Restart mode - prevent any write operations
            self.directory=restart
            self.restart=True
        else :
            self.directory=make_calc_dir(self.prefix,wdir)
            self.restart=False
        self.submited=False
        # override parameters with environment variables
        if 'ESPRESSO_PP_PATH' in os.environ:
            self.parameters['pseudo_dir'] = os.environ['ESPRESSO_PP_PATH']
        elif 'ESPRESSO_PSEUDO' in os.environ:
            self.parameters['pseudo_dir'] = os.environ['ESPRESSO_PSEUDO']
        if 'ASE_ESPRESSO_COMMAND' in os.environ:
            pw_cmd=os.environ['ASE_ESPRESSO_COMMAND']

    def copy(self):
        c=copy.deepcopy(self)
        c.directory=make_calc_dir(c.prefix,c.wdir)
        return c

    def build_command(self, prop=None, params=None):
        '''
        Build a propper QE command for the execution of a set of tasks.
        '''
        if prop is None:
            prop = ['energy']
        if params is None:
            params = {}
        # This is restarted calculation do not write anything.
        # We should not be even here. Return empty command.
        if self.restart : return ''
        
        params=params.copy()
        params.update({'prop': prop})
        cmd=''
        if set(['energy','stress','forces','bands','edos']) & set(prop):
            params.update({'iofile':'pw'})
            cmd+= self.pw_cmd % params
            cmd+='\n'
        if 'bands' in prop :
            params.update({'iofile':'bands','calc':'bands'})
            cmd+= self.pw_cmd % params
            cmd+='\n'
            # Move the result to the .bands directory 
            # Removing the previous .bands dir first
            cmd+=('rm -rf "' + os.path.join(params['outdir'],self.prefix+'.bands') + '"\n')
            cmd+=('cp -r ' + os.path.join(params['outdir'],self.prefix+'.save') + 
                ' ' + os.path.join(params['outdir'],self.prefix+'.bands') + ' \n')
        if 'edos' in prop :
            params.update({'iofile':'edos','calc':'nscf'})
            cmd+= self.pw_cmd % params
            cmd+='\n'
            cmd+= self.edos_cmd % params
            cmd+='\n'
            # Move the result to the .edos directory
            # Removing the previous .edos dir first
            cmd+=('rm -rf "' + os.path.join(params['outdir'],self.prefix+'.edos') + '"\n')
            cmd+=('cp -r ' + os.path.join(params['outdir'],self.prefix+'.save') + 
                ' ' + os.path.join(params['outdir'],self.prefix+'.edos') + ' \n')
        params.update({'calc':'scf'})
        if 'd2' in prop :
            cmd+= self.ph_cmd % params
            cmd+='\n'
            cmd+= self.q2r_cmd % params
            cmd+='\n'
        if 'frequencies' in prop :
            cmd+= self.matdyn_cmd % params
            cmd+='\n'
        if 'phdos' in prop :
            cmd+= self.phdos_cmd % params
            cmd+='\n'
        #print cmd
        return cmd
            

    def run_calculation(self, atoms, properties, system_changes):
        '''
        Hook for the more involved remote calculator to link into.
        Actually invokes the parent calculation routine.
        '''
        # This is restarted calculation do not write anything.
        # Nothing to do - just return.
        if self.restart : return
        
        FileIOCalculator.calculate(self, atoms, properties, system_changes)

    def calculate(self, atoms=None, properties=None,
                  system_changes=all_changes):
        '''
        Standard calculation interface
        '''
        if properties is None:
            properties = ['energy']
        # This is restarted calculation do not write anything.
        # Just read the results.
        if self.restart : 
            self.read_results()
            return
        
        self.command=self.build_command(properties,self.parameters)
        self.run_calculation(atoms, properties, system_changes)

    def write_input(self, atoms=None, properties=None, system_changes=None):
        """
        Write input file(s). This is called by FileIOCalculator just before 
        it executes the external command.
        """
        # This is restarted calculation do not write anything.
        if self.restart : return
        
        self.set(prefix=self.prefix)
        self.atoms2params()
        FileIOCalculator.write_input(self, atoms, properties, system_changes)
        
        if set(['energy','stress','forces','bands','edos']) & set(properties) :
            write_pw_in(self.directory, atoms, self.parameters)
        if 'bands' in properties :
            write_bands_in(self.directory, atoms, self.parameters)
        if 'edos' in properties:
            write_edos_in(self.directory, atoms, self.parameters)
        if 'd2' in properties :
            write_ph_in(self.directory, atoms, self.parameters)
            write_q2r_in(self.directory, atoms, self.parameters)
        if 'frequencies' in properties :
            write_matdyn_in(self.directory, atoms, self.parameters)
        if 'phdos' in properties:
            write_phdos_in(self.directory, atoms, self.parameters)
        



    def read_results(self):
        """Read energy, forces, ... from output file(s)."""
        
        fn=os.path.join(self.directory,'pw.out')
        xml_fn=os.path.join(self.directory,self.parameters['outdir'],self.prefix+'.save','data-file.xml')
        # Read the pan-ultimate line of the output file
        try: 
            ln=open(fn).readlines()[-2]
            if ln.find('JOB DONE.')>-1 :
                # Job is done we can read the output
                r=read_quantumespresso_textoutput(fn)
                self.results.update(r)
                self.results_xml={}
                self.results_xml.update(read_quantumespresso_xmloutput(xml_fn,'all'))
                self.results['energy']=r['etotal']*Rydberg
                try:
                  s=-array(r['stress'])* 1e-1 * ase.units.GPa
                  self.results['stress']=array([s[0, 0], s[1, 1], s[2, 2],
                                       s[1, 2], s[0, 2], s[0, 1]])
                except TypeError :
                  # stress is none - ignore.
                  pass
                        
                rk=self.results.keys()
                if 'cell' in rk :
                    try :
                        self.results['cell']=array(self.results['cell'])*self.results['alat']*Bohr
                    except TypeError:
                        # cell data not in order just clear the mess
                        del(self.results['cell'])
                if 'atoms_forces' in rk:
                    # Translate from the Ry/au of QE to the eV/A units of ASE
                    try :
                        self.results['forces']=array(self.results['atoms_forces'])*(ase.units.Rydberg/ase.units.Bohr)
                    except TypeError :
                        # atoms_forces are none - ignore.
                        pass

                try :
                    fn=os.path.join(self.directory,self.parameters['outdir'],self.prefix+'.bands','data-file.xml')
                    bands=read_quantumespresso_xmloutput(fn,'all')
                    # The result is in Hartree. How many energy units these italians will use?!
                    self.results['bands']=2*ase.units.Rydberg*array(bands['kpoints_eigenvalues']).T
                    self.results['bands_kpt']=array(bands['kpoints_coordinates'])
                except (IOError, KeyError, IndexError) :
                    # No bands - do nothing
                    pass

                try :
                    self.results['edos']=np.loadtxt(os.path.join(self.directory,self.prefix+'.edos')).T
                except IOError:
                    # No edos data. We assume there is none and silently
                    pass

                try :
                    self.results['phdos']=np.loadtxt(os.path.join(self.directory,self.prefix+'.dos')).T
                except IOError:
                    # No frequency data. We assume there is none and silently
                    pass

                try :
                    self.results['frequencies']=np.loadtxt(os.path.join(self.directory,self.prefix+'.freq.gp')).T
                except IOError:
                    # No frequency data. We assume there is none and silently
                    pass
            else :
                # Job not ready.
                raise CalcNotReadyError
        except (IOError, IndexError) :
            # Job not ready.
            raise CalcNotReadyError
        
    def clean(self):
        '''
        Celaning procedure. Non-functional for now.
        We need to decide which files should be left intact.
        '''
        # This is restarted calculation do not write anything.
        # And do not clean anything.
        if self.restart : return
        
        # Cleaning to be implemented 
        return
        
    def atoms2params(self):
        '''
            Populate parameters dictionary with data extracted from atoms.
            Return the updated dictionary.
        '''
        p=self.parameters
        return p



class RemoteQE(QuantumEspresso):
    '''
    Remote calculator class for Quantum Espresso.
    This class is only involved with the machanics of remotly executing
    the software and transporting the data. Any extensions concerning 
    actual calculation should be done in the QuantumEspresso class.
    '''
    
    # Queue system submit command
    qsub_tool='qsub'
    qstat_tool='qstat'
    
    qsub_cmd='cd %(rdir)s ; %(qsub_tool)s -N %(title)s -l procs=%(procs)d ./run-pw.pbs'

    # Remote execution command
    remote_exec_cmd='ssh %(user)s@%(host)s "%(command)s"'

    # If you cannot mount the data directory into your system it is best 
    # to use the rsync command to transfer the results back into the system.

    # Command for copying the data out to the computing system
    copy_out_cmd='rsync -a "%(ldir)s" "%(user)s@%(host)s:%(rdir)s"'
    # Command for copying the data in after the calculation
    copy_in_cmd='rsync -a "%(user)s@%(host)s:%(rdir)s" "%(ldir)s"'

    # Template for the PBS batch job
    pbs_template=''

    # Command to check the state of the job
    pbs_check_cmd='''%(qstat_tool)s -f %(jobid)s |grep job_state |awk '{print $3}' '''
    
    # Access data
    host=''
    user=''
    
    # Location:
    # local working directory
    wdir='.'
    # Remote working directory relative to the home directory or absolute
    rdir='.'

    # Repetition timer (seconds) for checkin the state of the job.
    job_check_time=15
    
    def __init__(self, block=True, **kwargs):
        QuantumEspresso.__init__(self, user=self.user, host=self.host, 
                                wdir=self.wdir, rdir=self.rdir, **kwargs)
        self.jobid=None
        self.block=block

    def write_pbs_in(self,properties):
        fh=open(os.path.join(self.directory,'run-pw.pbs'),'w')

        fh.write(self.pbs_template % {
            'command': QuantumEspresso.build_command(self,prop=properties,
                                                        params=self.parameters)
            })        
        
        fh.close()

    def build_command(self,prop=None,params=None):
        if prop is None:
            prop = ['energy']
        if params is None:
            params = {}
        cmd=self.qsub_cmd % {
            'qsub_tool': self.qsub_tool,
            'qstat_tool': self.qstat_tool,
            'title': self.label,
            'procs': self.parameters['procs'],
            'rdir': os.path.join(self.parameters['rdir'],os.path.split(self.directory)[-1])
        }
        cmd=self.remote_exec_cmd % {
                'command': cmd,
                'user': self.parameters['user'],
                'host': self.parameters['host']
         }
        return cmd

    def write_input(self, atoms=None, properties=None, system_changes=all_changes):
        '''Write input file(s).'''
        if properties is None:
            properties = ['energy']
        QuantumEspresso.write_input(self, atoms, properties, system_changes)
        self.write_pbs_in(properties)
        subprocess.call(self.copy_out_cmd % {
                            'ldir': self.directory,
                            'rdir': self.parameters['rdir'],
                            'user': self.parameters['user'],
                            'host': self.parameters['host']
                        }, shell=True)

    def job_ready(self):
        try :
            cmd=self.remote_exec_cmd % {
                'command': self.pbs_check_cmd % {
                                    'qsub_tool': self.qsub_tool,
                                    'qstat_tool': self.qstat_tool,
                                    'jobid':self.jobid
                                    },
                'user': self.parameters['user'],
                'host': self.parameters['host']
                }
            state=subprocess.check_output(cmd, shell=True).split()[-1]
        except (subprocess.CalledProcessError, IndexError) :
            # Unknown state. We assume it has finished and continue
            state='N'

        return not (state in ['Q','R'])


    def run_calculation(self, atoms=None, properties=None,
                            system_changes=all_changes):
        '''
        Internal calculation executor. We cannot use FileIOCalculator
        directly since we need to support remote execution.
        
        This calculator is different from others. 
        It prepares the directory, launches the remote process and
        raises the exception to signal that we need to come back for results
        when the job is finished.
        '''
        if properties is None:
            properties = ['energy']
        Calculator.calculate(self, atoms, properties, system_changes)
        self.write_input(self.atoms, properties, system_changes)
        if self.command is None:
            raise RuntimeError('Please configure RemoteQE calculator!')
        olddir = os.getcwd()
        errorcode=0
        try:
            os.chdir(self.directory)
            output = subprocess.check_output(self.command, shell=True)
            self.jobid=output.split()[0]
            self.submited=True
            #print "Job %s submitted. Waiting for it." % (self.jobid)
            # Waiting loop. To be removed.
        except subprocess.CalledProcessError as e:
            errorcode=e.returncode
        finally:
            os.chdir(olddir)
        
        if errorcode:
            raise RuntimeError('%s returned an error: %d' %
                               (self.name, errorcode))
        self.read_results()


    def read_results(self):
        """Read energy, forces, ... from output file(s)."""
        
        if self.submited:
            # The job has been submitted. Check the state.
            if not self.job_ready() :
                if self.block :
                    while not self.job_ready() :
                        time.sleep(self.job_check_time)
                else :
                    raise CalcNotReadyError
                    

            # Assume the calc finished. Copy the files back.
            subprocess.call(self.copy_in_cmd % {
                'ldir': self.wdir,
                'rdir': os.path.join(self.parameters['rdir'],os.path.split(self.directory)[-1]),
                'user': self.parameters['user'],
                'host': self.parameters['host']
            }, shell=True)
                
        
        fn=os.path.join(self.directory,'pw.out')
        # Read the pan-ultimate line of the output file
        try: 
            ln=open(fn).readlines()[-2]
            if ln.find('JOB DONE.')>-1 :
                # Job is done we can read the output
                r=read_quantumespresso_textoutput(fn)
                self.submited=False
                self.jobid=None
            else :
                # Job not ready.
                raise CalcNotReadyError
        except (IOError, IndexError) :
            # Job not ready.
            raise CalcNotReadyError
        
        # All is fine - really read the results
        QuantumEspresso.read_results(self)

    @classmethod
    def ParallelCalculate(cls,syslst,properties=None,system_changes=all_changes):
        '''
        Run a series of calculations in parallel using (implicitely) some 
        remote machine/cluster. The function returns the list of systems ready
        for the extraction of calculated properties.
        '''
        if properties is None:
                properties = ['energy']
        print 'Launching:',
        sys.stdout.flush()
        for n,s in enumerate(syslst):
            try :
                s.calc.block=False
                s.calc.calculate(atoms=s,properties=properties,system_changes=system_changes)
            except CalcNotReadyError:
                s.calc.block=True
            print n+1, 
            sys.stdout.flush()
        print
        print '     Done:',
        sys.stdout.flush()
        for n,s in enumerate(syslst):
            s.calc.read_results()
            print n+1, 
            sys.stdout.flush()
        print
        return syslst

