import sys
import numpy as np
from scipy.integrate import quad

# 分子構造例: H2
molecule = [
        [1, 0.0, 0.0, 0.0],
        [1, 0.0, 0.0, 1.4]  # Bohr単位で核間距離
    ]

mode = 'eh'
if len(sys.argv) > 1: mode = sys.argv[1]

HARTREE_TO_EV = 27.2114

# ==========================
# Slater型軌道 (簡易版)
# ==========================
def sto_norm_zeta(Z):
    """
    電荷Zに対するSlater軌道の指数ζと規格化定数Nを返す
    （教育用の簡易近似：ζ ≈ Z、有効電荷近似）
    """
    zeta = float(Z)
    N = (zeta**3 / np.pi)**0.5  # 1s STO の規格化定数
    return zeta, N

def sto_1s(r, zeta, N):
    return N * np.exp(-zeta * r)

# ==========================
# オーバーラップ積分
# ==========================
def overlap_numeric(zeta1, zeta2):
    """
    同じ中心の1s-STO同士の重なり積分（数値積分）
    """
    N1 = (zeta1**3 / np.pi)**0.5
    N2 = (zeta2**3 / np.pi)**0.5
    integrand = lambda r: sto_1s(r, zeta1, N1) * sto_1s(r, zeta2, N2) * r**2
    val, err = quad(integrand, 0, np.inf)
    return 4*np.pi*val

def overlap_analytic(zeta, R):
    """
    異なる中心の1s-STO間の既知の解析式
    S(R) = exp(-ζR) * (1 + ζR + (ζR)^2/3)
    """
    x = zeta * R
    return np.exp(-x) * (1 + x + x**2/3)

def overlap_mc(zeta, R, Nsample=100000, L=5.0):
    """
    モンテカルロ法で異中心1s軌道のオーバーラップを近似
    核A=(0,0,0), 核B=(0,0,R)
    """
    pts = np.random.uniform(-L, L, size=(Nsample, 3))
    rA = np.linalg.norm(pts, axis=1)
    rB = np.linalg.norm(pts - np.array([0,0,R]), axis=1)
    N = (zeta**3/np.pi)**0.5
    phiA = N * np.exp(-zeta * rA)
    phiB = N * np.exp(-zeta * rB)
    f = phiA * phiB
    volume = (2*L)**3
    return volume * np.mean(f)

# ==========================
# 運動・ポテンシャル・全エネルギー（解析式）
# ==========================
def kinetic_analytic(zeta):
    """ <T> = zeta^2 / 2  (Hartree) """
    return 0.5 * zeta**2

def potential_analytic(zeta):
    """ <V> = -zeta (Hartree) """
    return -zeta

def total_energy_analytic(zeta):
    """ E = <T> + <V> = zeta^2/2 - zeta (Hartree) """
    return kinetic_analytic(zeta) + potential_analytic(zeta)

# ==========================
# Extended Hückel Hamiltonian
# ==========================
def extended_huckel(molecule, K=1.75):
    """
    molecule: [[Z, x, y, z], ...] のリスト
    """
    N = len(molecule)
    H = np.zeros((N, N))
    S = np.zeros((N, N))

    # 対角項（イオン化ポテンシャル近似）
    onsite = []
    for atom in molecule:
        Z = atom[0]
        if Z == 1:   # H
            onsite.append(-13.6)
        elif Z == 6: # C
            onsite.append(-11.0)
        else:
            onsite.append(-10.0)  # 仮の値

    # 非対角項（オーバーラップを用いた近似）
    for i in range(N):
        H[i,i] = onsite[i]
        for j in range(i+1, N):
            Ri = np.array(molecule[i][1:])
            Rj = np.array(molecule[j][1:])
            R = np.linalg.norm(Ri - Rj)
            zeta_i, _ = sto_norm_zeta(molecule[i][0])
            zeta_j, _ = sto_norm_zeta(molecule[j][0])
            zeta = 0.5*(zeta_i+zeta_j)  # 平均ζ
            Sij = overlap_analytic(zeta, R)
            S[i,j] = Sij
            S[j,i] = Sij
            Hij = K * (H[i,i] + H[j,j]) / 2 * Sij
            H[i,j] = Hij
            H[j,i] = Hij

    eigvals, eigvecs = np.linalg.eigh(H)
    return eigvals, eigvecs, H, S

# ==========================
# メイン処理
# ==========================
def main(mode="STO"):

    if mode == "STO":
        for atom in molecule:
            zeta, N = sto_norm_zeta(atom[0])
            print(f"Z={atom[0]}: zeta={zeta:.2f}, Norm={N:.3f}")

    elif mode == "Sij":
        zeta1, _ = sto_norm_zeta(molecule[0][0])
        zeta2, _ = sto_norm_zeta(molecule[1][0])
        R = np.linalg.norm(np.array(molecule[0][1:]) - np.array(molecule[1][1:]))
        S_num = overlap_numeric(zeta1, zeta2)
        S_ana = overlap_analytic(zeta1, R)
        S_mc  = overlap_mc(zeta1, R, Nsample=200000)
        print(f"Numeric overlap (same center)        = {S_num:.4f}")
        print(f"Analytic overlap (different centers) = {S_ana:.4f}")
        print(f"Monte Carlo overlap (different centers) = {S_mc:.4f}")

    elif mode == "energy":
        # 例として H 1s (ζ=1)
        zeta = 1.0
        T = kinetic_analytic(zeta) * HARTREE_TO_EV
        V = potential_analytic(zeta) * HARTREE_TO_EV
        E = total_energy_analytic(zeta) * HARTREE_TO_EV
        print(f"Kinetic <T> (eV):    {T:.6f}")
        print(f"Potential <V> (eV):  {V:.6f}")
        print(f"Total E (eV):        {E:.6f}")

    elif mode == "eh":
        eigvals, eigvecs, H, S = extended_huckel(molecule)
        print("Eigenvalues (eV):", eigvals)
        print("Eigenvectors:\n", eigvecs)

if __name__ == "__main__":
    main(mode=mode)