"""A calculator based on the OpenKIM project."""
from __future__ import print_function

from ase.calculators.calculator import Calculator
from asap3 import _asap
from asap3.Internal.ListOfElements import ListOfElements
from asap3.Internal.BuiltinPotentials import get_cell_heights, smallest_pbc_direction
from ase.data import atomic_numbers, chemical_symbols
from ase.utils import basestring
import numpy as np
from copy import copy

OpenKIMinfo = _asap.OpenKIMinfo

kim_file_template = """#
# Autogenerated KIM specification file.
#
# Generated by the Asap module OpenKIMcalculator.
#

KIM_API_Version := 1.6.0

#  UNITS
Unit_length      := A
Unit_energy      := eV
Unit_charge      := e
Unit_temperature := K
Unit_time        := ps

SUPPORTED_ATOM/PARTICLES_TYPES:
# Symbol/name               Type                    code
%(symbols)s

CONVENTIONS:
# Name                      Type
ZeroBasedLists              flag
%(neighboraccess)s
%(neighbormethods)s

MODEL_INPUT:
# Name                      Type         Unit                Shape              Requirements
numberOfParticles           integer      none                []
numberOfSpecies             integer      none                []
particleSpecies             integer      none                [numberOfParticles]
coordinates                 double       length              [numberOfParticles,3]
get_neigh                   method       none                []
neighObject                 pointer      none                []
numberContributingParticles integer      none                []
boxSideLengths              double       length              [3]

MODEL_OUTPUT:
# Name                      Type         Unit                Shape              Requirements
destroy                     method       none                []
compute                     method       none                []
cutoff                      double       length              []
%(energy)s
%(forces)s
%(particleEnergy)s
%(virial)s
%(particleVirial)s
"""

element_template = "%-2s                          spec                    %i\n"
output_template = "%-27s double       %-19s [%s]"
nb_template = "%-27s flag\n"

class OpenKIMcalculator(_asap.OpenKIMcalculator):
    """A calculator interfacing to the OpenKIM models.
    
    Parameters:
    
    name: The long name of the KIM Model.
    
    Optional parameters:
    
    atoms:  If set, set_atoms is called immediately.  Default: no atoms set.
    
    allowed:  List of OpenKIM neighbor list keywords, only these will be considered when 
              matching the model.  Default: ALl neighbor list modes are allowed.
              
    allow_mi_opbc:  If set to False, minimum-image orthogonal periodic boundary condition
                    methods are not considered when mathcing the neighbor lists.  Useful
                    if the cell is expected to become skewed during the simulation.

    access:  Set to "iter" or "loca" to restrict the neigbor list access method to
             iterator mode or locator mode, respectively.  Default: not restricted.
             
    stress:  Set to False to refrain from calculate the global virial, even if the model
             supports it.  Default:  True (calculate global virial / stress).
             
    stresses: As above, but for atomic virials / stresses.
    
    verbose:  Set to True to print additional info during neigborlist matching.
    """
    def __init__(self, name, atoms=None, allowed=None, allow_mi_opbc=True, access=None,
                 stress=True, stresses=True, verbose=False):
        _asap.OpenKIMcalculator.__init__(self)
        self.name = name
        self.atoms = None
        self.allowed = allowed
        self.allow_mi_opbc = allow_mi_opbc
        self.access = access
        self.support_stress = stress
        self.support_stresses = stresses
        self.verbose = verbose
        if atoms is not None:
            self.set_atoms(atoms)
            
    def set_atoms(self, atoms, *args):
        """Set the atoms belonging to this calculator.
        
        The function set_atoms(atoms) is defined in C++, and just calls this
        Python function.  This ensures that this Python function is called
        regardless of whether set_atoms() is called from C++ or Python.
        """
        self.kim_spec = self.get_kim_spec(atoms)
        #print self.kim_spec
        self._initialize(self.kim_spec, self.name)
        # Inform the calculator about which quantities should be allocated and calculated.
        for k,v in self.supported_calculations.items():
            self.please_allocate(k, v)
        # Inform the calculator about the translation from atomic numbers to particle types.
        supported_elements = [atomic_numbers[sym] for sym in self.get_supported_types()]
        self.z_to_typecode = {}
        elements = ListOfElements(atoms)
        for e in elements:
            if e not in supported_elements:
                raise RuntimeError("The OpenKIM model '%s' does not support element Z=%i (%s)."
                                   % (self.name, e, chemical_symbols[e]))
            code = self.get_type_code(chemical_symbols[e])
            if code < 0:
                raise ValueError("Negative KIM type codes not supported - what do they mean?")
            self.z_to_typecode[e] = code
            if self.verbose:
                print("Translation: Z = %i -> id = %i" % (e, code))
        self.set_translation(self.z_to_typecode)
        # Find out which neighborlist type we got.  Check sanity and if ImageAtoms should be used.
        nblistmethod = self.get_NBC_method()
        pbc = atoms.get_pbc()
        cutoff = self.get_cutoff()
        if self.verbose:
            print("Neighbor list method: %s;  cutoff: %.3f A;  pbc: %s" % (nblistmethod, cutoff, pbc))
        if nblistmethod == 'CLUSTER':
            if not np.array_equal(pbc, (False, False, False)):
                raise RuntimeError("OpenKIM chose CLUSTER, but PBC are "+str(pbc))
        elif nblistmethod == 'NEIGH_RVEC_H' or nblistmethod == 'NEIGH_RVEC_F':
            if smallest_pbc_direction(atoms) < 3 * cutoff:
                if self.verbose:
                    print("Disabling minimum image convention for", nblistmethod)
                self._use_imageatoms()
        elif nblistmethod == "NEIGH_PURE_H" or nblistmethod == "NEIGH_PURE_F":
            if self.verbose:
                print("Activating Image atoms for", nblistmethod)
            self._use_imageatoms()
        elif nblistmethod == "MI_OPBC_H" or nblistmethod == "MI_OPBC_F":
            if not np.array_equal(pbc, (True, True, True)):
                raise RuntimeError("OpenKIM chose %s, but PBC are %s" (nblistmethod, str(pbc)))
            if smallest_pbc_direction(atoms) < 2 * cutoff:
                raise RuntimeError("OpenKIM chose %s, but system is too small." % (nblistmethod,))
        _asap.OpenKIMcalculator.set_atoms(self, atoms, *args)
        
    def get_kim_spec(self, atoms):
        """Return the KIM specification string.
        
        Side effects:
        self.supported_calculations (dictionary) will indicate which quantities
        can be calculated by the Model.
        self.z_to_typecode (dictionary) gives the translation from the 
        atomic numbers to the type codes of the model.
        """
        kim_variables = {'symbols': '', 'neighbormethods': ''}
        elements = ListOfElements(atoms)
        info = OpenKIMinfo(self.name)
        supported_elements = [atomic_numbers[sym] for sym in info.get_supported_types()]
        self.z_to_typecode = {}
        for e in elements:
            if e not in supported_elements:
                raise RuntimeError("The OpenKIM model '%s' does not support element Z=%i (%s)."
                                   % (self.name, e, chemical_symbols[e]))
            kim_variables['symbols'] += element_template % (chemical_symbols[e], e)
        test_methods = ('energy', 'particleEnergy', 'forces', 'virial', 'particleVirial')
        self.supported_calculations = {}
        for method in test_methods:
            self.supported_calculations[method] = info.get_API_index(method) > 0
        if not self.support_stress:
            self.supported_calculations['virial'] = False
        if not self.support_stresses:
            self.supported_calculations['particleVirial'] = False
        if self.supported_calculations['energy']:
            kim_variables['energy'] = output_template % ('energy', 'energy', '')
        else:
            kim_variables['energy'] = ''
        if self.supported_calculations['particleEnergy']:
            kim_variables['particleEnergy'] = output_template % ('particleEnergy', 'energy', 'numberOfParticles')
        else:
            kim_variables['particleEnergy'] = ''
        if self.supported_calculations['forces']:
            kim_variables['forces'] = output_template % ('forces', 'force', 'numberOfParticles,3')
        else:
            kim_variables['forces'] = ''
        if self.support_stress and self.supported_calculations['virial']:
            kim_variables['virial'] = output_template % ('virial', 'energy', '6')
        else:
            kim_variables['virial'] = ''
        if self.support_stress and self.supported_calculations['particleVirial']:
            kim_variables['particleVirial'] = output_template % ('particleVirial', 'energy', 'numberOfParticles,6')
        else:
            kim_variables['particleVirial'] = ''
        nb_methods = self.find_neighbor_methods(atoms)
        for nbm in nb_methods:
            kim_variables['neighbormethods'] += nb_template % (nbm,)
        if self.access == None:
            kim_variables['neighboraccess'] = "Neigh_IterAccess            flag\nNeigh_LocaAccess            flag"
        elif self.access.lower() == 'iter':
            if self.verbose:
                print("Using Neigh_IterAccess only")
            kim_variables['neighboraccess'] = "Neigh_IterAccess            flag"
        elif self.access.lower() == 'loca':
            if self.verbose:
                print("Using Neigh_LocaAccess only")
            kim_variables['neighboraccess'] = "Neigh_LocaAccess            flag"
        else:
            raise ValueError("Illegal value for neighborlist access method: " + str(self.access))
        return kim_file_template % kim_variables          
        
    def find_neighbor_methods(self, atoms):
        """Find which neighbor methods are relevant for this atom type
        
        Parameters:
        atoms: 
        The atoms object - its cell and PBC are examined.
        
        allowed (optional, default=None): 
        The list of neighborlist methods being considered, in prioritized order.  
        The default (None) means all methods, in the order NEIGH_RVEC_H, 
        NEIGH_PURE_H, NEIGH_RVEC_F, NEIGH_PURE_F, MI_OPBC_H, MI_OPBC_F, 
        CLUSTER.
        
        allow_mi_opbc (optional, default=True):
        If set to False, the methods MI_OPBC_H/F will not be considered.  Normally,
        they are considered if there are full periodic boundary conditions and the
        unit cell is orthorhombic.  However, the unit cell may change during a
        simulation, in those cases preventing these boundary conditions may be
        a good idea.
        """
        if not self.allowed:
            allowed = ["NEIGH_RVEC_H", "NEIGH_PURE_H", "NEIGH_RVEC_F", 
                       "NEIGH_PURE_F", "MI_OPBC_H", "MI_OPBC_F",
                       "CLUSTER"]
        elif isinstance(self.allowed, basestring):
            allowed = [self.allowed,]
        else:
            allowed = copy(self.allowed)  # Will be modified
        remove = []
        uc = atoms.get_cell()
        diagonal = True
        for i in range(3):
            for j in range(3):
                if i != j and np.abs(uc[i,j]) > 1e-15:
                    diagonal = False  
        if not (self.allow_mi_opbc and atoms.get_pbc().all() and diagonal):
            # Minimum Image Orthogonal Periodic Boundary Conditions
            # are not allowed
            remove.extend(["MI_OPBC_H", "MI_OPBC_F"])
        if  atoms.get_pbc().any():
            # Cluster method is not allowed
            remove.append("CLUSTER")
        for rem in remove:
            if rem in allowed:
                allowed.remove(rem)
        if self.verbose:
            print("Allowed PBC:", allowed)
        return allowed
    
    def get_parameters(self, kind='free', integers=[]):
        """Get the parameters of the model as a dictionary.
        
        The optional parameter 'kind' specifies which kind of parameters you get.
        It can take three values:
        'free' (the default).  You only get the "free" parameters, i.e. those that
            can be changed.
        'fixed': You only get the fixed parameters.
        'all': You get both free and fixed parameters.

        The optional argument 'integers' is a list of parameter names that are
        defined as integers in the parameter file.  This should be detected
        automatically once the KIM API specification is updated.
        
        WARNING: Currently using any other value than 'free' is dangerous, as the
        fixed parameter may be of unexpected types.  Once type-checking of parameters
        appear in the KIM API, accessing the fixed parameters will become safe, and 
        the default will change to 'all'.
        
        NOTE: The prefix PARAM_FREE_ or PARAM_FIXED_ is removed from the parameter 
        names.  If kind=='all', free parameters may shadow fixed parameters of the 
        same name.
        """
        if kind not in ['free', 'fixed', 'all']:
            raise ValueError("get_parameters: kind must be 'free', 'fixed' or 'all'.")
        result = {}
        if kind == 'fixed' or kind == 'all':
            prf = 'PARAM_FIXED_'
            for name in self.get_fixed_parameter_names():
                assert name.startswith(prf)
                shortname = name[len(prf):]
                if shortname in integers or name in integers:
                    value = self._get_parameter(name, is_integer=True)
                else:
                    value = self._get_parameter(name)
                result[shortname] = value
        if kind == 'free' or kind == 'all':
            prf = 'PARAM_FREE_'
            for name in self.get_free_parameter_names():
                assert name.startswith(prf)
                shortname = name[len(prf):]
                if shortname in integers or name in integers:
                    value = self._get_parameter(name, is_integer=True)
                else:
                    value = self._get_parameter(name)
                result[shortname] = value
        return result
    
    def set_parameters(self, **kwargs):
        """Set the free parameters of an OpenKIM model.
        
        Examples (these three are equivalent):
        
        model.set_parameters(cutoff = 3.7)
        
        model.set_parameters(PARAM_FREE_cutoff = 3.7)
        
        pars = {'cutoff': 3.7}
        model.set_parameters(**pars)
        
        Note that prefixing the parameter name with PARAM_FREE_ is optional.
        
        If a bad parameter specification is given, and this function fails,
        the model will be left in an undefined state.  That can be fixed by
        calling this function again, possibly without parameters.
        """
        prefix = 'PARAM_FREE_'
        allowed = self.get_free_parameter_names()
        for key, value in kwargs.items():
            if key.startswith(prefix):
                realkey = key
            else:
                realkey = prefix+key
            if realkey not in allowed:
                raise TypeError("Got an unexpected parameter name '{}'".format(key))
            if isinstance(value, float):
                self._set_parameter_scalar(realkey, value)
            elif isinstance(value, np.ndarray):
                self._set_parameter_array(realkey, value)
            else:
                raise TypeError("Unsupported value type - must be float or numpy array.")
        self._reinit()
    
if __name__ == '__main__':
    from ase.lattice.cubic import FaceCenteredCubic
    atoms = FaceCenteredCubic(size=(10,10,10), symbol='Cu')
    print("Creating calculator")
    pot = OpenKIMcalculator('EMT_Asap_Standard_AlAgAuCuNiPdPt__MO_118428466217_000')
    print("Setting atoms")
    atoms.set_calculator(pot)
    print("Calculating energy")
    print(atoms.get_potential_energy())
    print(atoms.get_forces()[10:])
    print(atoms.get_stress())

