# translated from cal_STO.pl by ChatGPT
# modified by Tomoya Suzuki

import sys


a0 = 5.29177e-11

AtomicNumber = 49
Charge = 3
unoccupied = True

argv = sys.argv
nargs = len(argv)
print("nargs=", nargs)
if nargs > 1:
    AtomicNumber = int(argv[1])
if nargs > 2:
    Charge = int(argv[2])
if nargs > 3:
    unoccupied = int(argv[3])



def calSTO(AtomicNumber = 20, Charge = 0, unoccupied = False):
    print(f"Atomic Number: {AtomicNumber}")
    print(f"Charge       : {Charge}")

    nElectron = AtomicNumber - Charge
    print(f"nElectron    : {nElectron}")
    print(f"calculate LUMO?: {unoccupied}")

    ResCharge = nElectron

    def subtract_charge(ResCharge, ne_orbital):
        if ResCharge > ne_orbital:
            n = ne_orbital
            ResCharge -= ne_orbital
        else:
            n = ResCharge
            ResCharge = 0.0
        return n, ResCharge

    #各軌道の電子数をカウント
    n1s, ResCharge = subtract_charge(ResCharge, 2.0)
    n2s, ResCharge = subtract_charge(ResCharge, 2.0)
    n2p, ResCharge = subtract_charge(ResCharge, 6.0)
    n3s, ResCharge = subtract_charge(ResCharge, 2.0)
    n3p, ResCharge = subtract_charge(ResCharge, 6.0)

    n3d = 0.0
    n4s = 0.0
    if Charge == 0.0:
        if ResCharge > 2.0:
            n4s = 2.0
            ResCharge -= 2.0
        else:
            n4s = ResCharge
            ResCharge = 0.0
        if ResCharge > 10.0:
            n3d = 10.0
            ResCharge -= 10.0
        else:
            n3d = ResCharge
            ResCharge = 0.0
    else:
        if ResCharge > 10.0:
            n3d = 10.0
            ResCharge -= 10.0
        else:
            n3d = ResCharge
            ResCharge = 0.0
        if ResCharge > 2.0:
            n4s = 2.0
            ResCharge -= 2.0
        else:
            n4s = ResCharge
            ResCharge = 0.0

    n4p, ResCharge = subtract_charge(ResCharge, 6.0)
    n4d, ResCharge = subtract_charge(ResCharge, 10.0)
    n5s, ResCharge = subtract_charge(ResCharge, 2.0)
    n5p, Rescharge = subtract_charge(ResCharge, 6.0)

    nTotal = n1s + n2s + n2p + n3s + n3p + n3d + n4s + n4p + n4d + n5s + n5p
    print(f"Total: {nTotal}")
    print(f" 1s: {n1s}")
    print(f" 2s: {n2s}")
    print(f" 2p: {n2p}")
    print(f" 3s: {n3s}")
    print(f" 3p: {n3p}")
    print(f" 3d: {n3d}")
    print(f" 4s: {n4s}")
    print(f" 4p: {n4p}")
    print(f" 4d: {n4d}")
    print(f" 5s: {n5s}")
    print(f" 5p: {n5p}")
    if ResCharge > 0.0:
        print(f"Error: Ne={AtomicNumber} - {Charge} is not supported.")
        exit()
    if nTotal != nElectron:
        print(f"Error: Inconsistent Ne (Ntotal={nTotal} != Ne={nElectron}).")
        exit()
    

    #空の軌道の電子を見たい場合。LUMO軌道に電子を入れて計算させる
    if unoccupied == True:
        print(f"search unoccupied orbital")
        orbital_group = [n1s, n2s, n2p, n3s, n3p, n3d, n4s, n4p, n4d, n5s, n5p]
        for num_orb in range(len(orbital_group)):
            if orbital_group[num_orb] == 0.0:
                orbital_group[num_orb] = 1.0
                break
        n1s, n2s, n2p, n3s, n3p, n3d, n4s, n4p, n4d, n5s, n5p = orbital_group
#        print("orbital_group=", orbital_group)

    #最外殻電子の主量子数を取得
    nGroup = [n1s, n2s + n2p, n3s + n3p + n3d, n4s + n4p + n4d, n5s + n5p]
    nMax = 0
    for i in range(len(nGroup)):
        if nGroup[i] > 0.0:
            nMax = i

    print()
    print(f"Max n: {nMax+1}")
    nStar = [1, 2, 3, 3.7, 4.0, 4.2]
    nStar = nStar[nMax]
    print(f"n*   : {nStar}")

    Is_d = 0
    if (n3d > 0.0 and n4s == 0.0 and unoccupied == False) or (n4d > 0.0 and n5s == 0.0 and unoccupied == False):
        Is_d = 1

    #遮蔽定数Sを計算
    stotal = 0.0
    
    # (iii) 考えている電子より主量子数nの大きい電子からのsへの寄与はゼロとする。
    # (iv) 考えている電子の属する群の他の電子からはそれぞれ0.35とする。ただし、1s群の場合は0.30とする。
    s = 0.0
    n = nGroup[nMax] - 1.0 #考えている電子の分引いてる
    k = 0.30 if nMax == 0 else 0.35
    if nGroup[nMax] > 1.0:
        s = n * k
    print(f"s(outmost)  = {n} * {k} = {s}")
    stotal += s

    # (v) 考えている電子がsまたはp電子なら、それよりnの一つ小さい電子からは0.85ずつ
    # ただし、d電子ならば、それより内側の電子からの寄与はすべて1.00とする。
    
    s = 0.0
    n = nGroup[nMax - 1]
    k = 1.0 if Is_d else 0.85
    if nMax >= 1 and nGroup[nMax - 1] > 1.0:
        s = n * k
    print(f"s(n(max)-1) = {n} * {k} = {s}")
    stotal += s

    # もっと内側の電子からは1.00ずつ寄与する。
    k = 1.0
    n = 0.0
    for i in range(nMax - 2, -1, -1):
        n += nGroup[i]
    s = n * k
    print(f"s(inner)    = {n} * {k} = {s}")
    stotal += s
    print(f"s(total)    = {stotal}")
    
    #有効核電荷計算
    zStar = AtomicNumber - stotal
    print(f"Target atom: Z={AtomicNumber}  q={Charge} nElectron={nElectron}")
    print(f"LUMO?: {unoccupied}")
    print(f"Z*   : {zStar}")

    #slater 半径計算
    Rmax = (nStar * nStar) / zStar * a0 * 1.0e10
    n1 = nStar - 1
    print(f"R(r) = N * r^{n1} * exp[-({zStar}/{nStar} a0) * r]")
    print(f"     a0={a0} angstrom")
    print(f"Rmax = (a0/Z*)(n*)^2 = {Rmax} (A)")
    print()

#実行
#空の軌道を見たい場合はunoccupied=Trueに
calSTO(AtomicNumber = AtomicNumber, Charge = Charge, unoccupied = unoccupied)

input("\nPress ENTER to terminate>>")