#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
XC functional, the interface to xcfun (https://github.com/dftlibs/xcfun)
U. Ekstrom et al, J. Chem. Theory Comput., 6, 1971
'''

import copy
import ctypes
import math
import numpy
from pyscf import lib

_itrf = lib.load_library('libxcfun_itrf')

XC = XC_CODES = {
'SLATERX'       :  0,  # Slater LDA exchange
'PW86X'         :  1,  # PW86 exchange
'VWN3C'         :  2,  # VWN3 LDA Correlation functional
'VWN5C'         :  3,  # VWN5 LDA Correlation functional
'PBEC'          :  4,  # PBE correlation functional
'PBEX'          :  5,  # PBE Exchange Functional
'BECKEX'        :  6,  # Becke 88 exchange
'BECKECORRX'    :  7,  # Becke 88 exchange correction
'BECKESRX'      :  8,  # Short range Becke 88 exchange
'BECKECAMX'     :  9,  # CAM Becke 88 exchange
'BRX'           : 10,  # Becke-Roussells exchange with jp dependence
'BRC'           : 11,  # Becke-Roussells correlation with jp dependence
'BRXC'          : 12,  # Becke-Roussells correlation with jp dependence
'LDAERFX'       : 13,  # Short-range spin-dependent LDA exchange functional
'LDAERFC'       : 14,  # Short-range spin-dependent LDA correlation functional
'LDAERFC_JT'    : 15,  # Short-range spin-unpolarized LDA correlation functional
'LYPC'          : 16,  # LYP correlation
'OPTX'          : 17,  # OPTX Handy & Cohen exchange
'OPTXCORR'      : 18,  # OPTX Handy & Cohen exchange -- correction part only
'REVPBEX'       : 19,  # Revised PBE Exchange Functional
'RPBEX'         : 20,  # RPBE Exchange Functional
'SPBEC'         : 21,  # sPBE correlation functional
'VWN_PBEC'      : 22,  # PBE correlation functional using VWN LDA correlation.
'KTX'           : 23,  # KT exchange GGA correction
'TFK'           : 24,  # Thomas-Fermi Kinetic Energy Functional
'TW'            : 25,  # von Weizsacker Kinetic Energy Functional
'PW91X'         : 26,  # Perdew-Wang 1991 GGA Exchange Functional
'PW91K'         : 27,  # PW91 GGA Kinetic Energy Functional
'PW92C'         : 28,  # PW92 LDA correlation
'M05X'          : 29,  # M05 exchange
'M05X2X'        : 30,  # M05-2X exchange
'M06X'          : 31,  # M06 exchange
'M06X2X'        : 32,  # M06-2X exchange
'M06LX'         : 33,  # M06-L exchange
'M06HFX'        : 34,  # M06-HF exchange
'M05X2C'        : 35,  # M05-2X Correlation
'M05C'          : 36,  # M05 Correlation
'M06C'          : 37,  # M06 Correlation
'M06HFC'        : 38,  # M06-HF Correlation
'M06LC'         : 39,  # M06-L Correlation
'M06X2C'        : 40,  # M06-2X Correlation
'TPSSC'         : 41,  # TPSS original correlation functional
'TPSSX'         : 42,  # TPSS original exchange functional
'REVTPSSC'      : 43,  # Revised TPSS correlation functional
'REVTPSSX'      : 44,  # Reviewed TPSS exchange functional
'PZ81C'         : 45,  # PZ81 LDA correlation
'P86C'          : 46,  # P86C GGA correlation
'P86CORRC'      : 47,  # P86C GGA correlation
'BTK'           : 48,  # Borgoo-Tozer TS
'VWK'           : 49,  # von Weizsaecker kinetic energy
'B97X'          : 50,  # B97 exchange
'B97C'          : 51,  # B97 correlation
'B97_1X'        : 52,  # B97-1 exchange
'B97_1C'        : 53,  # B97-1 correlation
'B97_2X'        : 54,  # B97-2 exchange
'B97_2C'        : 55,  # B97-2 correlation
'CSC'           : 56,  # Colle-Salvetti correlation functional
'APBEC'         : 57,  # APBE correlation functional.
'APBEX'         : 58,  # APBE Exchange Functional
'ZVPBESOLC'     : 59,  # zvPBEsol correlation Functional
'BLOCX'         : 60,  # BLOC exchange functional
'PBEINTC'       : 61,  # PBEint correlation Functional
'PBEINTX'       : 62,  # PBEint Exchange Functional
'PBELOCC'       : 63,  # PBEloc correlation functional.
'PBESOLX'       : 64,  # PBEsol Exchange Functional
'TPSSLOCC'      : 65,  # TPSSloc correlation functional
'ZVPBEINTC'     : 66,  # zvPBEint correlation Functional
'PW91C'         : 67,  # PW91 Correlation
#
# alias
#
'SLATER'        : 0,  # SLATERX
'LDA'           : 0,  # SLATERX
'VWN'           : 3,  # VWN5C
'VWN5'          : 3,  # VWN5C
'VWN3'          : 2,  # VWN3C
'SVWN'          : 'SLATERX + VWN5',
'B88'           : 6,  # BECKECORRX
'LYP'           : 16,
'P86'           : 46,
'M052XX'        : 30,  # M05-2X exchange
'M062XX'        : 32,  # M06-2X exchange
'M052XC'        : 35,  # M05-2X Correlation
'M062XC'        : 40,  # M06-2X Correlation
'BLYP'          : 'B88 + LYP',
'BP86'          : 'B88 + P86',  # Becke-Perdew 1986
'BPW91'         : 'B88 + PW91C',
'BPW92'         : 'B88 + PW92C',
'OLYP'          : '2.4832*SLATER - 1.43169*OPTX + LYP',  # CPL, 341, 319
'KT1'           : 'SLATERX - 0.006*KTX',  # Keal-Tozer 1, JCP, 119, 3015
'KT2XC'         : '1.07173*SLATER - .006*KTX + 0.576727*VWN5',  # Keal-Tozer 2, JCP, 119, 3015
'KT3XC'         : 'SLATERX*1.092 + KTX*-0.004 + OPTXCORR*-0.925452 + LYPC*0.864409',  # Keal-Tozer 3, JCP, 121, 5654
# == '2.021452*SLATER - .004*KTX - .925452*OPTX + .864409*LYP',
'PBE0'          : '.25*HF + .75*PBEX + PBEC',  # Perdew-Burke-Ernzerhof, JCP, 110, 6158
'PBE1PBE'       : 'PBE0',
'PBEH'          : 'PBE0',
'B3P86'         : '.2*HF + .08*SLATER + .72*B88 + .81*P86C + .19*VWN5C',
'B3P86G'        : '.2*HF + .08*SLATER + .72*B88 + .81*P86C + .19*VWN3C',
'B3PW91'        : '.2*HF + .08*SLATER + .72*B88 + .81*PW91C + .19*VWN5C',
'B3PW91G'       : '.2*HF + .08*SLATER + .72*B88 + .81*PW91C + .19*VWN3C',
# Note, use VWN5 for B3LYP. It is different to the libxc default B3LYP
'B3LYP'         : 'B3LYP5',
'B3LYP5'        : '.2*HF + .08*SLATER + .72*B88 + .81*LYP + .19*VWN5C',
'B3LYPG'        : '.2*HF + .08*SLATER + .72*B88 + .81*LYP + .19*VWN3C', # B3LYP-VWN3 used by Gaussian and libxc
#'O3LYP'         : '.1161*HF + .9262*SLATER + .8133*OPTXCORR + .81*LYP + .19*VWN5C',  # Mol. Phys. 99 607
#'O3LYPG'        : '.1161*HF + .9262*SLATER + .8133*OPTXCORR + .81*LYP + .19*VWN3C',
# libxc implementation as below, see also discussion in https://gitlab.com/libxc/libxc/issues/47
#'O3LYP'         : '.1161*HF + .9262*SLATER + 1.164393477*OPTXCORR + .81*LYP + .19*VWN5C', #1.164393477 = .8133*1.43169
#'O3LYPG'        : '.1161*HF + .9262*SLATER + 1.164393477*OPTXCORR + .81*LYP + .19*VWN3C',
'O3LYP'         : '.1161*HF + 0.071006917*SLATER + .8133*OPTX, .81*LYP + .19*VWN5',  # libxc implementation
'X3LYP'         : '.218*HF + .073*SLATER + 0.542385*B88 + .166615*PW91X + .871*LYP + .129*VWN5C',  # Xu, PNAS, 101, 2673
'X3LYPG'        : '.218*HF + .073*SLATER + 0.542385*B88 + .166615*PW91X + .871*LYP + .129*VWN3C',
# Range-separated-hybrid functional: (alpha+beta)*SR_HF(0.33) + alpha*LR_HF(0.33)
# Note default mu of xcfun is 0.4. It can cause discrepancy for CAMB3LYP
'CAMB3LYP'      : '0.19*SR_HF(0.33) + 0.65*LR_HF(0.33) + BECKECAMX + VWN5C*0.19 + LYPC*0.81',
'CAM_B3LYP'     : 'CAMB3LYP',
'LDAERF'        : 'LDAERFX + LDAERFC',  # Short-range exchange and correlation LDA functional
'B97XC'         : 'B97X + B97C + HF*0.1943',
'B97_1XC'       : 'B97_1X + B97_1C + HF*0.21',
'B97_2XC'       : 'B97_2X + B97_2C + HF*0.21',
'M05XC'         : '.28*HF + .72*M05X + M05C',
'TPSSH'         : '0.1*HF + 0.9*TPSSX + TPSSC',
'TF'            : 'TFK',
}

# Some XC functionals have conventional name, like M06-L means M06-L for X
# functional and M06-L for C functional, PBE mean PBE-X plus PBE-C. If the
# conventional name was placed in the XC_CODES, it may lead to recursive
# reference when parsing the xc description.  These names (as exceptions of
# XC_CODES) are listed in XC_ALIAS below and they should be treated as a
# shortcut for XC functional.
XC_ALIAS = {
    # Conventional name : name in XC_CODES
    'BLYP'              : 'B88,LYP',
    'BP86'              : 'B88,P86',
    'PW91'              : 'PW91,PW91',
    'PBE'               : 'PBE,PBE',
    'REVPBE'            : 'REVPBE,PBE',
    'PBESOL'            : 'PBESOL,PBESOL',
#    'PKZB'              : 'PKZB,PKZB',
    'TPSS'              : 'TPSS,TPSS',
    'REVTPSS'           : 'REVTPSS,REVTPSS',
#    'SCAN'              : 'SCAN,SCAN',
#    'SOGGA'             : 'SOGGA,PBE',
    'BLOC'              : 'BLOC,TPSSLOC',
    'OLYP'              : 'OPTX,LYP',
    'RPBE'              : 'RPBE,PBE',
    'BPBE'              : 'B88,PBE',
#    'MPW91'             : 'MPW91,PW91',
    'HFLYP'             : 'HF,LYP',
#    'HFPW92'            : 'HF,PWMOD',
#    'SPW92'             : 'SLATER,PWMOD',
    'SVWN'              : 'SLATER,VWN',
#    'MS0'               : 'MS0,REGTPSS',
#    'MS1'               : 'MS1,REGTPSS',
#    'MS2'               : 'MS2,REGTPSS',
#    'MS2H'              : 'MS2H,REGTPSS',
#    'MVS'               : 'MVS,REGTPSS',
#    'MVSH'              : 'MVSH,REGTPSS',
#    'SOGGA11'           : 'SOGGA11,SOGGA11',
#    'SOGGA11-X'         : 'SOGGA11X,SOGGA11X',
    'KT1'               : 'KT1,VWN',
#    'DLDF'              : 'DLDF,DLDF',
#    'GAM'               : 'GAM,GAM',
    'M06-L'             : 'M06L,M06L',
#    'M11-L'             : 'M11L,M11L',
#    'MN12-L'            : 'MN12L,MN12L',
#    'MN15-L'            : 'MN15L,MN15L',
#    'N12'               : 'N12,N12',
#    'N12-SX'            : 'N12SX,N12SX',
#    'MN12-SX'           : 'MN12SX,MN12SX',
#    'MN15'              : 'MN15,MN15',
#    'MBEEF'             : 'MBEEF,PBESOL',
#    'SCAN0'             : 'SCAN0,SCAN',
#    'PBEOP'             : 'PBE,OPPBE',
#    'BOP'               : 'B88,OPB88',
}
XC_ALIAS.update([(key.replace('-',''), XC_ALIAS[key])
                 for key in XC_ALIAS if '-' in key])

LDA_IDS = set([0, 2, 3, 13, 14, 15, 24, 28, 45])
GGA_IDS = set([1, 4, 5, 6, 7, 8, 9, 16, 17, 18, 19, 20, 21, 22, 23, 25, 26,
               27, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 57, 58, 59, 61,
               62, 63, 64, 66, 67])
MGGA_IDS =set([10, 11, 12, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41,
               42, 43, 56, 60, 65])
HYB_XC = set(('PBE0'    , 'PBE1PBE' , 'B3PW91'  , 'B3P86'   , 'B3LYP'   ,
              'B3PW91G' , 'B3P86G'  , 'B3LYPG'  , 'O3LYP'   , 'CAMB3LYP',
              'B97XC'   , 'B97_1XC' , 'B97_2XC' , 'M05XC'   , 'TPSSH'   ,
              'HFLYP'))
RSH_XC = set(('CAMB3LYP',))
MAX_DERIV_ORDER = 3

VV10_DAT = {
    'B97M_V'    : [6.0, 0.01],
    'WB97M_V'   : [6.0, 0.01],
    'WB97X_V'   : [6.0, 0.01],
    'VV10'      : [5.9, 0.0093],
    'LC_VV10'   : [6.3, 0.0089],
}
VV10_XC = set(VV10_DAT.keys())

def xc_type(xc_code):
    if isinstance(xc_code, str):
        if is_nlc(xc_code):
            return 'NLC'
        hyb, fn_facs = parse_xc(xc_code)
    else:
        fn_facs = [(xc_code, 1)]  # mimic fn_facs
    if not fn_facs:
        return 'HF'
    elif all(xid in LDA_IDS for xid, val in fn_facs):
        return 'LDA'
    elif any(xid in MGGA_IDS for xid, val in fn_facs):
        return 'MGGA'
    else:
        # all((xid in GGA_IDS or xid in LDA_IDS for xid, val in fn_fns)):
        # include hybrid_xc
        return 'GGA'

def is_lda(xc_code):
    return xc_type(xc_code) == 'LDA'

def is_hybrid_xc(xc_code):
    if isinstance(xc_code, str):
        xc_code = xc_code.replace(' ','').upper()
        return ('HF' in xc_code or xc_code in HYB_XC or
                hybrid_coeff(xc_code) != 0)
    elif isinstance(xc_code, int):
        return False
    else:
        return any((is_hybrid_xc(x) for x in xc_code))

def is_meta_gga(xc_code):
    return xc_type(xc_code) == 'MGGA'

def is_gga(xc_code):
    return xc_type(xc_code) == 'GGA'

def is_nlc(xc_code):
    return xc_code.upper() in VV10_XC

def nlc_coeff(xc_code):
    '''Get NLC coefficients
    '''
    xc_code = xc_code.upper()
    if is_nlc(xc_code):
        return VV10_DAT[xc_code]
    else:
        return 0, 0

def rsh_coeff(xc_code):
    '''Get Range-separated-hybrid coefficients
    '''
    hyb, fn_facs = parse_xc(xc_code)
    hyb, alpha, omega = hyb
    beta = hyb - alpha
    return omega, alpha, beta

def max_deriv_order(xc_code):
    hyb, fn_facs = parse_xc(xc_code)
    return MAX_DERIV_ORDER

def test_deriv_order(xc_code, deriv, raise_error=False):
    support = deriv <= max_deriv_order(xc_code)
    if not support and raise_error:
        raise NotImplementedError('xcfun library does not support derivative '
                                  'order %d for  %s' % (deriv, xc_code))
    return support

def hybrid_coeff(xc_code, spin=0):
    if is_nlc(xc_code):
        return 0
    hyb, fn_facs = parse_xc(xc_code)
    return hyb[0]

def parse_xc_name(xc_name):
    fn_facs = parse_xc(xc_name)[1]
    return fn_facs[0][0], fn_facs[1][0]

def parse_xc(description):
    r'''Rules to input functional description:

    * The given functional description must be a one-line string.
    * The functional description is case-insensitive.
    * The functional description string has two parts, separated by ",".  The
      first part describes the exchange functional, the second is the correlation
      functional.

      - If "," was not in string, the entire string is considered as a
        compound XC functional (including both X and C functionals, such as b3lyp).
      - To input only X functional (without C functional), leave the second
        part blank. E.g. description='slater,' means pure LDA functional.
      - To neglect X functional (just apply C functional), leave the first
        part blank. E.g. description=',vwn' means pure VWN functional.
      - If compound XC functional is specified, no matter whehter it is in the
        X part (the string in front of comma) or the C part (the string behind
        comma), both X and C functionals of the compound XC functional will be
        used.

    * The functional name can be placed in arbitrary order.  Two name needs to
      be separated by operators "+" or "-".  Blank spaces are ignored.
      NOTE the parser only reads operators "+" "-" "*".  / is not in support.
    * A functional name can have at most one factor.  If the factor is not
      given, it is set to 1.  Compound functional can be scaled as a unit. For
      example '0.5*b3lyp' is equivalent to
      'HF*0.1 + .04*LDA + .36*B88, .405*LYP + .095*VWN'
    * String "HF" stands for exact exchange (HF K matrix).  Putting "HF" in
      correlation functional part is the same to putting "HF" in exchange
      part.
    * String "RSH" means range-separated operator. Its format is
      RSH(alpha; beta; omega).  Another way to input RSH is to use keywords
      SR_HF and LR_HF: "SR_HF(0.1) * alpha_plus_beta" and "LR_HF(0.1) *
      alpha" where the number in parenthesis is the value of omega.
    * Be careful with the convention on GGA functional, in which the LDA
      contribution has been included.
    '''
    hyb = [0, 0, 0]  # hybrid, alpha, omega
    if isinstance(description, int):
        return 0, [(description, 1.)]
    elif not isinstance(description, str): #isinstance(description, (tuple,list)):
        return parse_xc('%s,%s' % tuple(description))

    def assign_omega(omega, hyb_or_sr, lr=0):
        if hyb[2] == omega or omega == 0:
            hyb[0] += hyb_or_sr
            hyb[1] += lr
        elif hyb[2] == 0:
            hyb[0] += hyb_or_sr
            hyb[1] += lr
            hyb[2] = omega
        else:
            raise ValueError('Different values of omega found for RSH functionals')
    fn_facs = []
    def parse_token(token, suffix):
        if token:
            if token[0] == '-':
                sign = -1
                token = token[1:]
            else:
                sign = 1
            if '*' in token:
                fac, key = token.split('*')
                if fac[0].isalpha():
                    fac, key = key, fac
                fac = sign * float(fac)
            else:
                fac, key = sign, token

            if key[:3] == 'RSH':
# RSH(alpha; beta; omega): Range-separated-hybrid functional
                alpha, beta, omega = [float(x) for x in key[4:-1].split(';')]
                assign_omega(omega, fac*(alpha+beta), fac*alpha)
            elif key == 'HF':
                hyb[0] += fac
                hyb[1] += fac  # also add to LR_HF
            elif 'SR_HF' in key or 'SRHF' in key:
                if '(' in key:
                    omega = float(key.split('(')[1].split(')')[0])
                    assign_omega(omega, fac, 0)
                else:  # Assuming this omega the same to the existing omega
                    hyb[0] += fac
            elif 'LR_HF' in key:
                if '(' in key:
                    omega = float(key.split('(')[1].split(')')[0])
                    assign_omega(omega, 0, fac)
                else:
                    hyb[1] += fac  # == alpha
            elif key.isdigit():
                fn_facs.append((int(key), fac))
            else:
                if key in XC_CODES:
                    x_id = XC_CODES[key]
                elif key+suffix in XC_CODES:
                    x_id = XC_CODES[key+suffix]
                else:
                    raise KeyError('Unknown XC key %s' % key)
                if isinstance(x_id, str):
                    hyb1, fn_facs1 = parse_xc(x_id)
# Recursively scale the composed functional, to support e.g. '0.5*b3lyp'
                    if hyb1[0] != 0 or hyb1[1] != 0:
                        assign_omega(hyb1[2], hyb1[0]*fac, hyb1[1]*fac)
                    fn_facs.extend([(xid, c*fac) for xid, c in fn_facs1])
                elif x_id is None:
                    raise NotImplementedError(key)
                else:
                    fn_facs.append((x_id, fac))
    def remove_dup(fn_facs):
        fn_ids = []
        facs = []
        n = 0
        for key, val in fn_facs:
            if key in fn_ids:
                facs[fn_ids.index(key)] += val
            else:
                fn_ids.append(key)
                facs.append(val)
                n += 1
        return list(zip(fn_ids, facs))

    description = description.replace(' ','').upper()
    if description in XC_ALIAS:
        description = XC_ALIAS[description]

    if '-' in description:  # To handle e.g. M06-L
        for key in _NAME_WITH_DASH:
            if key in description:
                description = description.replace(key, _NAME_WITH_DASH[key])

    if ',' in description:
        x_code, c_code = description.split(',')
        for token in x_code.replace('-', '+-').split('+'):
            parse_token(token, 'X')
        for token in c_code.replace('-', '+-').split('+'):
            parse_token(token, 'C')
    else:
        for token in description.replace('-', '+-').split('+'):
            parse_token(token, 'XC')
    if hyb[2] == 0: # No omega is assigned. LR_HF is 0 for normal Coulomb operator
        hyb[1] = 0
    return hyb, remove_dup(fn_facs)

_NAME_WITH_DASH = {'SR-HF'  : 'SR_HF',
                   'LR-HF'  : 'LR_HF',
                   'M06-L'  : 'M06L',
                   'M05-2X' : 'M052X',
                   'M06-HF' : 'M06HF',
                   'M06-2X' : 'M062X',}


def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, verbose=None):
    r'''Interface to call xcfun library to evaluate XC functional, potential
    and functional derivatives.

    See also :func:`pyscf.dft.libxc.eval_xc`
    '''
    hyb, fn_facs = parse_xc(xc_code)
    return _eval_xc(fn_facs, rho, spin, relativity, deriv, verbose)

XC_D0 = 0
XC_D1 = 1
XC_D2 = 2
XC_D3 = 3
XC_D4 = 4

XC_D00 = 0
XC_D10 = 1
XC_D01 = 2
XC_D20 = 3
XC_D11 = 4
XC_D02 = 5
XC_D30 = 6
XC_D21 = 7
XC_D12 = 8
XC_D03 = 9
XC_D40 = 10
XC_D31 = 11
XC_D22 = 12
XC_D13 = 13
XC_D04 = 14

XC_D000 = 0
XC_D100 = 1
XC_D010 = 2
XC_D001 = 3
XC_D200 = 4
XC_D110 = 5
XC_D101 = 6
XC_D020 = 7
XC_D011 = 8
XC_D002 = 9
XC_D300 = 10
XC_D210 = 11
XC_D201 = 12
XC_D120 = 13
XC_D111 = 14
XC_D102 = 15
XC_D030 = 16
XC_D021 = 17
XC_D012 = 18
XC_D003 = 19
XC_D400 = 20
XC_D310 = 21
XC_D301 = 22
XC_D220 = 23
XC_D211 = 24
XC_D202 = 25
XC_D130 = 26
XC_D121 = 27
XC_D112 = 28
XC_D103 = 29
XC_D040 = 30
XC_D031 = 31
XC_D022 = 32
XC_D013 = 33
XC_D004 = 34

XC_D00000 = 0
XC_D10000 = 1
XC_D01000 = 2
XC_D00100 = 3
XC_D00010 = 4
XC_D00001 = 5
XC_D20000 = 6
XC_D11000 = 7
XC_D10100 = 8
XC_D10010 = 9
XC_D10001 = 10
XC_D02000 = 11
XC_D01100 = 12
XC_D01010 = 13
XC_D01001 = 14
XC_D00200 = 15
XC_D00110 = 16
XC_D00101 = 17
XC_D00020 = 18
XC_D00011 = 19
XC_D00002 = 20
XC_D30000 = 21
XC_D21000 = 22
XC_D20100 = 23
XC_D20010 = 24
XC_D20001 = 25
XC_D12000 = 26
XC_D11100 = 27
XC_D11010 = 28
XC_D11001 = 29
XC_D10200 = 30
XC_D10110 = 31
XC_D10101 = 32
XC_D10020 = 33
XC_D10011 = 34
XC_D10002 = 35
XC_D03000 = 36
XC_D02100 = 37
XC_D02010 = 38
XC_D02001 = 39
XC_D01200 = 40
XC_D01110 = 41
XC_D01101 = 42
XC_D01020 = 43
XC_D01011 = 44
XC_D01002 = 45
XC_D00300 = 46
XC_D00210 = 47
XC_D00201 = 48
XC_D00120 = 49
XC_D00111 = 50
XC_D00102 = 51
XC_D00030 = 52
XC_D00021 = 53
XC_D00012 = 54
XC_D00003 = 55
XC_D40000 = 56
XC_D31000 = 57
XC_D30100 = 58
XC_D30010 = 59
XC_D30001 = 60
XC_D22000 = 61
XC_D21100 = 62
XC_D21010 = 63
XC_D21001 = 64
XC_D20200 = 65
XC_D20110 = 66
XC_D20101 = 67
XC_D20020 = 68
XC_D20011 = 69
XC_D20002 = 70
XC_D13000 = 71
XC_D12100 = 72
XC_D12010 = 73
XC_D12001 = 74
XC_D11200 = 75
XC_D11110 = 76
XC_D11101 = 77
XC_D11020 = 78
XC_D11011 = 79
XC_D11002 = 80
XC_D10300 = 81
XC_D10210 = 82
XC_D10201 = 83
XC_D10120 = 84
XC_D10111 = 85
XC_D10102 = 86
XC_D10030 = 87
XC_D10021 = 88
XC_D10012 = 89
XC_D10003 = 90
XC_D04000 = 91
XC_D03100 = 92
XC_D03010 = 93
XC_D03001 = 94
XC_D02200 = 95
XC_D02110 = 96
XC_D02101 = 97
XC_D02020 = 98
XC_D02011 = 99
XC_D02002 = 100
XC_D01300 = 101
XC_D01210 = 102
XC_D01201 = 103
XC_D01120 = 104
XC_D01111 = 105
XC_D01102 = 106
XC_D01030 = 107
XC_D01021 = 108
XC_D01012 = 109
XC_D01003 = 110
XC_D00400 = 111
XC_D00310 = 112
XC_D00301 = 113
XC_D00220 = 114
XC_D00211 = 115
XC_D00202 = 116
XC_D00130 = 117
XC_D00121 = 118
XC_D00112 = 119
XC_D00103 = 120
XC_D00040 = 121
XC_D00031 = 122
XC_D00022 = 123
XC_D00013 = 124
XC_D00004 = 125

XC_D0000000 = 0
XC_D1000000 = 1
XC_D0100000 = 2
XC_D0010000 = 3
XC_D0001000 = 4
XC_D0000100 = 5
XC_D0000010 = 6
XC_D0000001 = 7
XC_D2000000 = 8
XC_D1100000 = 9
XC_D1010000 = 10
XC_D1001000 = 11
XC_D1000100 = 12
XC_D1000010 = 13
XC_D1000001 = 14
XC_D0200000 = 15
XC_D0110000 = 16
XC_D0101000 = 17
XC_D0100100 = 18
XC_D0100010 = 19
XC_D0100001 = 20
XC_D0020000 = 21
XC_D0011000 = 22
XC_D0010100 = 23
XC_D0010010 = 24
XC_D0010001 = 25
XC_D0002000 = 26
XC_D0001100 = 27
XC_D0001010 = 28
XC_D0001001 = 29
XC_D0000200 = 30
XC_D0000110 = 31
XC_D0000101 = 32
XC_D0000020 = 33
XC_D0000011 = 34
XC_D0000002 = 35
XC_D3000000 = 36
XC_D2100000 = 37
XC_D2010000 = 38
XC_D2001000 = 39
XC_D2000100 = 40
XC_D2000010 = 41
XC_D2000001 = 42
XC_D1200000 = 43
XC_D1110000 = 44
XC_D1101000 = 45
XC_D1100100 = 46
XC_D1100010 = 47
XC_D1100001 = 48
XC_D1020000 = 49
XC_D1011000 = 50
XC_D1010100 = 51
XC_D1010010 = 52
XC_D1010001 = 53
XC_D1002000 = 54
XC_D1001100 = 55
XC_D1001010 = 56
XC_D1001001 = 57
XC_D1000200 = 58
XC_D1000110 = 59
XC_D1000101 = 60
XC_D1000020 = 61
XC_D1000011 = 62
XC_D1000002 = 63
XC_D0300000 = 64
XC_D0210000 = 65
XC_D0201000 = 66
XC_D0200100 = 67
XC_D0200010 = 68
XC_D0200001 = 69
XC_D0120000 = 70
XC_D0111000 = 71
XC_D0110100 = 72
XC_D0110010 = 73
XC_D0110001 = 74
XC_D0102000 = 75
XC_D0101100 = 76
XC_D0101010 = 77
XC_D0101001 = 78
XC_D0100200 = 79
XC_D0100110 = 80
XC_D0100101 = 81
XC_D0100020 = 82
XC_D0100011 = 83
XC_D0100002 = 84
XC_D0030000 = 85
XC_D0021000 = 86
XC_D0020100 = 87
XC_D0020010 = 88
XC_D0020001 = 89
XC_D0012000 = 90
XC_D0011100 = 91
XC_D0011010 = 92
XC_D0011001 = 93
XC_D0010200 = 94
XC_D0010110 = 95
XC_D0010101 = 96
XC_D0010020 = 97
XC_D0010011 = 98
XC_D0010002 = 99
XC_D0003000 = 100
XC_D0002100 = 101
XC_D0002010 = 102
XC_D0002001 = 103
XC_D0001200 = 104
XC_D0001110 = 105
XC_D0001101 = 106
XC_D0001020 = 107
XC_D0001011 = 108
XC_D0001002 = 109
XC_D0000300 = 110
XC_D0000210 = 111
XC_D0000201 = 112
XC_D0000120 = 113
XC_D0000111 = 114
XC_D0000102 = 115
XC_D0000030 = 116
XC_D0000021 = 117
XC_D0000012 = 118
XC_D0000003 = 119

def _eval_xc(fn_facs, rho, spin=0, relativity=0, deriv=1, verbose=None):
    assert(deriv < 4)
    if spin == 0:
        rho_u = rho_d = numpy.asarray(rho, order='C')
    else:
        rho_u = numpy.asarray(rho[0], order='C')
        rho_d = numpy.asarray(rho[1], order='C')
    assert(rho_u.dtype == numpy.double)
    assert(rho_d.dtype == numpy.double)

    if rho_u.ndim == 2:
        ngrids = rho_u.shape[1]
    else:
        ngrids = len(rho_u)

    fn_ids = [x[0] for x in fn_facs]
    facs   = [x[1] for x in fn_facs]
    if all((is_lda(x) for x in fn_ids)):  # LDA
        if spin == 0:
            nvar = 1
        else:
            nvar = 2
    elif any((is_meta_gga(x) for x in fn_ids)):
        if spin == 0:
            nvar = 3
        else:
            nvar = 7
    else:  # GGA
        if spin == 0:
            nvar = 2
        else:
            nvar = 5
    outlen = (math.factorial(nvar+deriv) //
              (math.factorial(nvar) * math.factorial(deriv)))
    outbuf = numpy.empty((ngrids,outlen))

    n = len(fn_ids)
    _itrf.XCFUN_eval_xc(ctypes.c_int(n),
                        (ctypes.c_int*n)(*fn_ids), (ctypes.c_double*n)(*facs),
                        ctypes.c_int(spin),
                        ctypes.c_int(deriv), ctypes.c_int(ngrids),
                        rho_u.ctypes.data_as(ctypes.c_void_p),
                        rho_d.ctypes.data_as(ctypes.c_void_p),
                        outbuf.ctypes.data_as(ctypes.c_void_p))

    outbuf = outbuf.T
    exc = outbuf[0]
    vxc = fxc = kxc = None
    if nvar == 1:
        if deriv > 0:
            vxc = (outbuf[1], None, None, None)
        if deriv > 1:
            fxc = (outbuf[2],) + (None,)*9
        if deriv > 2:
            kxc = (outbuf[3], None, None, None)
    elif nvar == 2:
        if spin == 0:  # GGA
            if deriv > 0:
                vxc = (outbuf[1], outbuf[2], None, None)
            if deriv > 1:
                fxc = (outbuf[3], outbuf[4], outbuf[5],) + (None,)*7
            if deriv > 2:
                kxc = outbuf[6:10]
        else:  # LDA
            if deriv > 0:
                vxc = (outbuf[1:3].T, None, None, None)
            if deriv > 1:
                fxc = (outbuf[3:6].T,) + (None,)*9
            if deriv > 2:
                kxc = (outbuf[6:10].T, None, None, None)
    elif nvar == 5:
        if deriv > 0:
            vxc = (outbuf[1:3].T, outbuf[3:6].T, None, None)
        if deriv > 1:
            fxc = (outbuf[[XC_D20000,XC_D11000,XC_D02000]].T,
                   outbuf[[XC_D10100,XC_D10010,XC_D10001,
                           XC_D01100,XC_D01010,XC_D01001]].T,
                   outbuf[[XC_D00200,XC_D00110,XC_D00101,XC_D00020,XC_D00011,XC_D00002]].T) + (None,)*7
        if deriv > 2:
            kxc = (outbuf[[XC_D30000,XC_D21000,XC_D12000,XC_D03000]].T,
                   outbuf[[XC_D20100,XC_D20010,XC_D20001,
                           XC_D11100,XC_D11010,XC_D11001,
                           XC_D02100,XC_D02010,XC_D02001]].T,
                   outbuf[[XC_D10200,XC_D10110,XC_D10101,XC_D10020,XC_D10011,XC_D10002,
                           XC_D01200,XC_D01110,XC_D01101,XC_D01020,XC_D01011,XC_D01002]].T,
                   outbuf[[XC_D00300,XC_D00210,XC_D00201,XC_D00120,XC_D00111,
                           XC_D00102,XC_D00030,XC_D00021,XC_D00012,XC_D00003]].T)
# MGGA/MLGGA: Note the MLGGA interface are not implemented. MGGA only needs 3
# input arguments.  To make the interface compatible with libxc, treat MGGA as
# MLGGA
    elif nvar == 3:
        if deriv > 0:
            vxc = (outbuf[1], outbuf[2], numpy.zeros_like(outbuf[1]), outbuf[3])
        if deriv > 1:
            fxc = (outbuf[XC_D200], outbuf[XC_D110], outbuf[XC_D020],
                   None, outbuf[XC_D002], None, outbuf[XC_D101], None, None, outbuf[XC_D011])
        if deriv > 2:
            kxc = (outbuf[XC_D300], outbuf[XC_D210], outbuf[XC_D120], outbuf[XC_D030],
                   outbuf[XC_D201], outbuf[XC_D111], outbuf[XC_D102],
                   outbuf[XC_D021], outbuf[XC_D012], outbuf[XC_D003])
    elif nvar == 7:
        if deriv > 0:
            vxc = (outbuf[1:3].T, outbuf[3:6].T, None, outbuf[6:8].T)
        if deriv > 1:
            fxc = (outbuf[[XC_D2000000,XC_D1100000,XC_D0200000]].T,
                   outbuf[[XC_D1010000,XC_D1001000,XC_D1000100,
                           XC_D0110000,XC_D0101000,XC_D0100100]].T,
                   outbuf[[XC_D0020000,XC_D0011000,XC_D0010100,
                           XC_D0002000,XC_D0001100,XC_D0000200]].T,
                   None,
                   outbuf[[XC_D0000020,XC_D0000011,XC_D0000002]].T,
                   None,
                   outbuf[[XC_D1000010,XC_D1000001,XC_D0100010,XC_D0100001]].T,
                   None, None,
                   outbuf[[XC_D0010010,XC_D0010001,XC_D0001010,XC_D0001001,
                           XC_D0000110,XC_D0000101]].T)
        if deriv > 2:
            kxc = (outbuf[[XC_D3000000,XC_D2100000,XC_D1200000,XC_D0300000]].T,
                   outbuf[[XC_D2010000,XC_D2001000,XC_D2000100,
                           XC_D1110000,XC_D1101000,XC_D1100100,
                           XC_D0210000,XC_D0201000,XC_D0200100]].T,
                   outbuf[[XC_D1020000,XC_D1011000,XC_D1010100,XC_D1002000,XC_D1001100,XC_D1000200,
                           XC_D0120000,XC_D0111000,XC_D0110100,XC_D0102000,XC_D0101100,XC_D0100200]].T,
                   outbuf[[XC_D0030000,XC_D0021000,XC_D0020100,XC_D0012000,XC_D0011100,
                           XC_D0010200,XC_D0003000,XC_D0002100,XC_D0001200,XC_D0000300]].T,
                   outbuf[[XC_D2000010,XC_D2000001,XC_D1100010,XC_D1100001,XC_D0200010,XC_D0200001]].T,
                   outbuf[[XC_D1010010,XC_D1010001,XC_D1001010,XC_D1001001,XC_D1000110,XC_D1000101,
                           XC_D0110010,XC_D0110001,XC_D0101010,XC_D0101001,XC_D0100110,XC_D0100101]].T,
                   outbuf[[XC_D1000020,XC_D1000011,XC_D1000002,XC_D0100020,XC_D0100011,XC_D0100002]].T,
                   outbuf[[XC_D0020010,XC_D0020001,XC_D0011010,XC_D0011001,XC_D0010110,XC_D0010101,
                           XC_D0002010,XC_D0002001,XC_D0001110,XC_D0001101,XC_D0000210,XC_D0000201]].T,
                   outbuf[[XC_D0010020,XC_D0010011,XC_D0010002,
                           XC_D0001020,XC_D0001011,XC_D0001002,
                           XC_D0000120,XC_D0000111,XC_D0000102]].T,
                   outbuf[[XC_D0000030,XC_D0000021,XC_D0000012,XC_D0000003]].T)
    return exc, vxc, fxc, kxc


def define_xc_(ni, description, xctype='LDA', hyb=0, rsh=(0,0,0)):
    '''Define XC functional.  See also :func:`eval_xc` for the rules of input description.

    Args:
        ni : an instance of :class:`NumInt`

        description : str
            A string to describe the linear combination of different XC functionals.
            The X and C functional are separated by comma like '.8*LDA+.2*B86,VWN'.
            If "HF" was appeared in the string, it stands for the exact exchange.

    Kwargs:
        xctype : str
            'LDA' or 'GGA' or 'MGGA'
        hyb : float
            hybrid functional coefficient
        rsh : float
            coefficients for range-separated hybrid functional

    Examples:

    >>> mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz')
    >>> mf = dft.RKS(mol)
    >>> define_xc_(mf._numint, '.2*HF + .08*LDA + .72*B88, .81*LYP + .19*VWN')
    >>> mf.kernel()
    -76.3783361189611
    >>> define_xc_(mf._numint, 'LDA*.08 + .72*B88 + .2*HF, .81*LYP + .19*VWN')
    >>> mf.kernel()
    -76.3783361189611
    >>> def eval_xc(xc_code, rho, *args, **kwargs):
    ...     exc = 0.01 * rho**2
    ...     vrho = 0.01 * 2 * rho
    ...     vxc = (vrho, None, None, None)
    ...     fxc = None  # 2nd order functional derivative
    ...     kxc = None  # 3rd order functional derivative
    ...     return exc, vxc, fxc, kxc
    >>> define_xc_(mf._numint, eval_xc, xctype='LDA')
    >>> mf.kernel()
    48.8525211046668
    '''
    if isinstance(description, str):
        ni.eval_xc = lambda xc_code, rho, *args, **kwargs: \
                eval_xc(description, rho, *args, **kwargs)
        ni.hybrid_coeff = lambda *args, **kwargs: hybrid_coeff(description)
        ni._xc_type = lambda *args: xc_type(description)

    elif callable(description):
        ni.eval_xc = description
        ni.hybrid_coeff = lambda *args, **kwargs: hyb
        ni.rsh_coeff = lambda *args, **kwargs: rsh
        ni._xc_type = lambda *args: xctype
    else:
        raise ValueError('Unknown description %s' % description)
    return ni

def define_xc(ni, description, xctype='LDA', hyb=0, rsh=(0,0,0)):
    return define_xc_(copy.copy(ni), description, xctype, hyb, rsh)
define_xc.__doc__ = define_xc_.__doc__


if __name__ == '__main__':
    from pyscf import gto, dft
    mol = gto.M(
        atom = [
        ["O" , (0. , 0.     , 0.)],
        [1   , (0. , -0.757 , 0.587)],
        [1   , (0. , 0.757  , 0.587)] ],
        basis = '6311g',)
    mf = dft.RKS(mol)
    mf._numint.libxc = dft.xcfun
    print(mf.kernel() - -75.8503877483363)

    mf.xc = 'b88,lyp'
    print(mf.kernel() - -76.3969707800463)

    mf.xc = 'b3lyp'
    print(mf.kernel() - -76.3777689410509)
