#!/usr/bin/env python3


"""
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.

Usage:
      qp_set_frozen_core [-q|--query] [(-l|-s|-u|--large|--small|--unset)] EZFIO_DIR

Options:
    -q --query   Prints in the standard output the number of frozen MOs
    -l --large   Use a large core
    -s --small   Use a small core
    -u --unset   Unset frozen core


Default numbers of frozen electrons:

      ========== ========= ======= ======= 
        Range     Default   Small   Large  
      ========== ========= ======= ======= 
       H  -> He      0        0        0   
       Li -> Be      0        0        2   
       B  -> Ne      2        2        2   
       Na -> Mg      2        2       10   
       Al -> Ar     10        2       10   
       K  -> Ca     10       10       18   
       Sc -> Zn     10       10       18   
       Ga -> Kr     18       10       18   
       Rb -> Sr     18       18       36   
       Y  -> Cd     18       18       36   
       In -> Xe     36       18       36   
       Cs -> Ba     36       36       54   
       La -> Hg     36       36       54   
       Tl -> Rn     54       36       54   
       Fr -> Ra     54       54       86   
       Ac -> Cn     54       54       86   
       Nh -> Og     86       54       86
      ========== ========= ======= ======= 


"""

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


    if not do_pseudo:

        if arguments["--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

        elif arguments["--small"]:
                if   charge <=   4: pass
                elif charge <=  18: n_frozen += 1
                elif charge <=  36: n_frozen += 5
                elif charge <=  54: n_frozen += 9
                elif charge <=  86: n_frozen += 18
                elif charge <= 118: n_frozen += 27
        elif arguments["--unset"]:

              n_frozen = 0
        else:  # default                                                                      
            for charge in ezfio.nuclei_nucl_charge:
                if   charge <=   4: pass
                elif charge <=  12: n_frozen += 1
                elif charge <=  30: n_frozen += 5
                elif charge <=  48: n_frozen += 9
                elif charge <=  80: n_frozen += 18
                elif charge <= 112: 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)