#!/usr/bin/env python2


"""
Automatically finds n, the number of core electrons. Calls qp_set_mo_class
setting all MOs as Active, except the n/2 first ones which are set as Core.
If pseudo-potentials are used, all the MOs are set as Active.

For elements on the right of the periodic table, qp_set_frozen_core will work
as expected. But for elements on the left, a small core will be chosen. For
example, a Carbon atom will have 2 core electrons, but a Lithium atom will have
zero.


Usage:
      qp_set_frozen_core [-q|--query] [-l|--large] EZFIO_DIR

Options:
    -q --query   Prints in the standard output the number of frozen MOs
    -l --large   Use a large core

"""

import os
import sys
import os.path

try:
    import qp_path
except ImportError:
    print "source .quantum_package.rc"
    raise

from docopt import docopt
from ezfio import ezfio


def main(arguments):
    """Main function"""

    filename = arguments["EZFIO_DIR"]
    ezfio.set_filename(filename)

    n_frozen = 0
    try:
        do_pseudo = ezfio.pseudo_do_pseudo
    except:
        do_pseudo = False

    large = 0
    small = 1

    size = small
    if arguments["--large"]:
        size = large

    if not do_pseudo:

        if size == large:
            for charge in ezfio.nuclei_nucl_charge:
		if charge <= 2:
		    pass 
		elif charge <= 10:
		    n_frozen += 1
		elif charge <= 18:
		    n_frozen += 5
		elif charge <= 36:
		    n_frozen += 9
		elif charge <= 54:
		    n_frozen += 18
		elif charge <= 86:
		    n_frozen += 27
		elif charge <= 118:
		    n_frozen += 43

        if size == small:

            for charge in ezfio.nuclei_nucl_charge:
                if charge < 5:
                    pass
                elif charge < 13:
                    n_frozen += 1
                elif charge < 31:
                    n_frozen += 5
                elif charge < 49:
                    n_frozen += 9
                elif charge < 81:
                    n_frozen += 18
                elif charge < 113:
                    n_frozen += 27

    mo_num = ezfio.mo_basis_mo_num

    if arguments["--query"]:
        print n_frozen
        sys.exit(0)

    if n_frozen == 0:
        os.system("""qp_set_mo_class -a "[1-%d]" %s""" %
                  (mo_num, filename))
    else:
        os.system("""qp_set_mo_class -c "[1-%d]" -a "[%d-%d]" %s""" %
                  (n_frozen, n_frozen+1, mo_num, filename))



if __name__ == '__main__':
    ARGUMENTS = docopt(__doc__)
    main(ARGUMENTS)