import os
import re
import sys

def A2au(r):
    if isinstance(r, list):
        for i in range(len(r)):
            r[i]=r[i]/0.52917725
        return r
    else:
        return r/0.52917725

def au2A(r):
    if isinstance(r, list):
        for i in range(len(r)):
            r[i]=r[i]*0.52917725
        return r
    else:
        return r*0.52917725


def getNElectrons(symbol, xmlName):
    """ 
    adapted frompseudopotential_section function by
    Author: "Kittithat (Mick) Krongchon" <pikkienvd3@gmail.com>
    in cif2crystal.

    Returns the integer of electrons in the pseudoatom

    Args:
        xymbos (str): The symbol of the element
        xmlName (str): The name of the XML pseudopotential and basis set database.
    Returns:
        int: electrons in pseudoatom
    """
    tree = ElementTree()
    try:
        tree.parse(xmlName)
    except Exception, inst:
        sys.exit("Error opening PP file: {}".format(inst))
    element = tree.find('./Pseudopotential[@symbol="{}"]'.format(symbol))
    effCoreCharge = int(element.find('./Effective_core_charge').text)

    return effCoreCharge

def getNOrbitals(symbol, xmlNames, basisNames):
    """
    adapted from basis_modified function by
    Author: "Kittithat (Mick) Krongchon" <pikkienvd3@gmail.com>
    in cif2crystal.

    Returns number of orbitals in the basis

    Args:
       symbol (str): The symbol of the element
       xmlName (str): List of  names of the XML pseudopotential and basis set database.
       basisName (str): List of names of the basis set (VTZ, ...)

    Returns:
       int: number of orbitals
    """
    nOrbs=0
    for i in range(len(xmlNames)):
        xmlName=xmlNames[i]
        basisName=basisNames[i]
        tree = ElementTree()
        try:
            tree.parse(xmlName)
        except Exception, inst:
            sys.exit("Error opening PP file: {}".format(inst))
        element = tree.find('./Pseudopotential[@symbol="{}"]'.format(symbol))
        if element==None:
            sys.exit("No basis found for {}".format(symbol))

        #use G functions only on explicit request
        useG=False
        if basisName[-2:]=="_G":
            basisName=basisName[:-2]
            useG=True
        if (useG):
            angsPass=['h', 'i']
        else:
            angsPass=['g', 'h', 'i']


        #changes 2.9.2015 (also added use of G on request)
        basis=element.find('./Basis-set[@name="{}"]'.format(basisName))
        if basis==None:
                sys.exit("No basis {} found for {}".format(basisName, symbol))

        #get contractions from xml file
        for contraction in basis.findall('./Contraction'):
            angular = contraction.get('Angular_momentum')
            if angular=='s':
                nOrbs+=1
            elif angular=='p':
                nOrbs+=3
            elif angular=='d':
                nOrbs+=5
            elif angular=='f':
                nOrbs+=7
            elif angular in angsPass:
                pass
            #don't count orbitals h and higher, they are not included in the basis!
            #only count g on request
            elif angular in 'g':
                nOrbs+=9
            else:
                 sys.exit("unknown angular momentum")
    return nOrbs