#!/usr/bin/env python3
import os
import sys
import argparse
import numpy as np

from pymatgen.core.structure import Structure
from pymatgen.core.sites import PeriodicSite
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer


# --- グローバル定義 ---
tkProg_Root = os.environ.get('tkProg_Root')
if not tkProg_Root:
    print("\nError: Environment variable tkProg_Root must be specified\n")
    sys.exit(1)

db_path = os.path.join(tkProg_Root, 'tkdb', 'Databases', 'bvparm2020.cif')


def terminate():
    input("Press ENTER to terminate>> ")

def read_bv_params(path):
    bv_params = {}
    if not os.path.exists(path):
        print(f"Error: File not found at [{path}]")
        return None
    with open(path, 'r', encoding='utf-8') as f:
        lines = f.readlines()

    reading = False
    for line in lines:
        line = line.strip()
        if not line:
            continue
        if line == '_valence_param_details':
            reading = True
            continue
        if reading:
            if line.startswith('#'):
                break
            parts = line.split()
            if len(parts) < 7:
                print(f"Warning: skipping malformed line: {line}")
                continue
            atom1, z1_str, atom2, z2_str, r0_str, b_str, ref = parts[:7]
            details = " ".join(parts[7:]) if len(parts) > 7 else ''
            try:
                z1 = int(z1_str); z2 = int(z2_str)
                r0 = float(r0_str); B = float(b_str)
            except ValueError as e:
                print(f"Warning: conversion error in line: {line}  ({e})")
                continue
            key = f"{atom1}:{atom2}"
            if key in bv_params:
#                print(f"Warning: duplicate key '{key}', skipping")
                continue
            bv_params[key] = {
                'atom1': atom1, 'Z1': z1,
                'atom2': atom2, 'Z2': z2,
                'r0': r0, 'B': B,
                'ref': ref, 'details': details
            }
    return bv_params

def main():
    parser = argparse.ArgumentParser(
        description="Calculate Bond Valence Sums from CIF (custom BV parameters)")
    parser.add_argument("cif_file", help="Input CIF file")
    parser.add_argument("max_r", nargs='?', type=float, default=2.7,
                        help="Cutoff distance for neighbors (Å)")
    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)

    bv_params = read_bv_params(db_path)
    if bv_params is None:
        sys.exit(1)

    try:
        structure = Structure.from_file(args.cif_file)
    except Exception as e:
        print(f"Error reading CIF: {e}")
        sys.exit(1)

# --- 対称性解析 ---
    sga = SpacegroupAnalyzer(structure, symprec=1e-3)  # 必要なら symprec を調整
    symm_struct = sga.get_symmetrized_structure()

    isite = 0
#    for i, site1 in enumerate(structure.sites):
    for eq_group in symm_struct.equivalent_sites:
# symm_struct.equivalent_sites は [ [site1a, site1b,...], [site2a,...], ... ] のリスト
# representative は各グループの最初のサイトでOK
        site1 = eq_group[0]   # 代表サイト
        i = structure.sites.index(site1)
        if hasattr(site1, "specie"):
            name1 = site1.specie.symbol
        else:
            print(f"\nError to obtain site1.specie.symbol: pymatgen may fail to read partial occupancy\n")
            terminate()
        
        try:
            charge1 = site1.specie.oxi_state
            if charge1 is None:
                charge1 = 0
        except AttributeError:
            charge1 = 0
        print(f"{isite:2d}: {name1}{charge1 if charge1!=0 else ''}")
        isite += 1
        bvs = 0.0

        for neighbor in structure.get_neighbors(site1, args.max_r):
            site2 = neighbor
            d = neighbor.nn_distance
            if d == 0.0:
                continue

            if hasattr(site2, "specie"):
                name2 = site2.specie.symbol
            else:
                print(f"\nError to obtain site2.specie.symbol: pymatgen may fail to read partial occupancy\n")
                terminate()

            try:
                charge2 = site2.specie.oxi_state
                if charge2 is None:
                    charge2 = 0
            except AttributeError:
                charge2 = 0

            neighbor_idx = neighbor.index

            fwd = f"{name1}:{name2}"
            bwd = f"{name2}:{name1}"
            params = bv_params.get(fwd) or bv_params.get(bwd)
            if not params:
                print(f"  Error: no BV params for {fwd} or {bwd}, skipping")
                continue

            r0 = params['r0']; B = params['B']
            if B == 0:
                print(f"  Warning: B=0 for key '{fwd}', skipping")
                continue

            bv = np.exp((r0 - d) / B)
            bvs += bv
            j = structure.sites.index(site2) if site2 in structure.sites else -1
            print(f"  {j:2d}: {name2}{charge2 if charge2!=0 else '':<4}"
                  f" r0={r0:.3f} B={B:.2f} d={d:.3f} bv={bv:.3f}"
                  f" [{params['ref']}]")

        print(f"    Bond Valence Sum: {bvs:.3f}\n")

    terminate()

if __name__ == "__main__":
    main()