import os
import sys
import argparse
import pprint

from pymatgen.core.structure import Structure
from pymatgen.analysis.bond_valence import BVAnalyzer
from pymatgen.analysis.bond_valence import calculate_bv_sum
import pymatgen.analysis.bond_valence as bv_mod


# BVAnalyzer
#著者：N. E. Brese & M. O’Keeffe
#タイトル：Bond-valence parameters for solids
#雑誌：Acta Crystallographica Section B: Structural Science, Crystal Engineering and Materials
#巻 (Volume)：47
#号 (Issue)：2
#ページ：192–197
#発行年：1991年（April 1991）
#DOI：10.1107/S0108768190011041
# O’Keefe & Brese の文献に基づくパラメータセット（元素ごとにr,cを決め, B=0.31）を使用
# C:\Users\tkami\AppData\Local\Programs\Python\Python312\Lib\site-packages\pymatgen\analysis\bvparam_1991.yaml
# 例: Zn,Oのr,c から、r0を計算
# r0 = r_zn + r_o - (r_zn * r_o * (sqrt(c_zn) - sqrt(c_o))**2) / (c_zn * r_zn + c_o * r_o)


def main():
    parser = argparse.ArgumentParser(
        description="Calculate Bond Valence Sums from a CIF file using pymatgen"
    )
    parser.add_argument("cif_file", help="Path to the CIF file")
    parser.add_argument("--max_r", type=float, default=2.8,
                        help="Maximum neighbor distance for BVAnalyzer (Å)")
    args = parser.parse_args()

    if not os.path.isfile(args.cif_file):
        print(f"Error: CIF file not found: {args.cif_file}")
        sys.exit(1)

    structure = Structure.from_file(args.cif_file)
    bva = BVAnalyzer(max_radius=args.max_r)

    print()
    print(f"cif_file: {args.cif_file}")
    print(f"max_r: {args.max_r}")

    print()
    print("=== BVAnalyzer settings ===")
    print(f"symm_tol                 = {bva.symm_tol}")
    print(f"max_radius               = {bva.max_radius}")
    print(f"distance_scale_factor    = {bva.dist_scale_factor}")
    print(f"max_permutations         = {bva.max_permutations}")
    print(f"charge_neutrality_tol    = {bva.charge_neutrality_tolerance}")
#    print(f"forbidden_species        = {bva.forbidden_species}")
    param_dir = os.path.dirname(bv_mod.__file__)
    param_file = os.path.join(param_dir, "bvparam_1991.yaml")
    print("\n=== Default parameter file ===")
    print(param_file)
    print("\n=== Example: Zn–O parameters ===")
# BV_PARAMS は {element: {neighbor: {"r":…, "c":…}, …}, …}
    pprint.pprint(bv_mod.BV_PARAMS.get("Zn", {}).get("O"))

    print()
    print("Calculate BVS:")
    for site in structure.sites:
        nn = structure.get_neighbors(site, args.max_r)
        bvs = calculate_bv_sum(site, nn, scale_factor=bva.dist_scale_factor)
        print(f"site {site.specie.symbol} nn_count={len(nn)} BVS={bvs:.3f}")

    print()
    print("\n=== calculate_bv_sum で距離を変えてみる ===")
    for rcut in [0.0, 1.5, 2.0, 2.5, 3.0]:
        bva = BVAnalyzer(max_radius=rcut)
        try:
            structure = bva.get_oxi_state_decorated_structure(structure)
        except:
            break

        # 各サイトごとに nearest neighbors を取り、BVS を直接計算
        for site in structure.sites:
            nn = structure.get_neighbors(site, rcut)
            bvs = calculate_bv_sum(site, nn, scale_factor=bva.dist_scale_factor)
            print(f"rcut={rcut:3.1f} Å, site={site.specie.symbol}, nn={len(nn)}, BVS={bvs:.3f}")
        print("-" * 40)
    
    print("BVAnalyzer result:")
    try:
        valences = bva.get_valences(structure)
    except:
        valences = None
    if valences:
        for i, (site, bv_sum) in enumerate(zip(structure.sites, valences)):
            symbol = site.specie.symbol
            coords = site.frac_coords
            print(f"{i:2d}: {symbol:2s} at {coords}  →  BVS = {bv_sum:.3f}")

    input("\nPress ENTER to terminate>>\n")
    
if __name__ == "__main__":
    main()