"""
为指定案例添加次紧邻键 (k2, k=100, r0=1.41421356)。

用法:  python add_k2.py case16
       python add_k2.py case16 case17 case18
       python add_k2.py --all
"""

import os
import sys

def add_k2(case_dir):
    coord_path = os.path.join(case_dir, "input", "coord.txt")
    conn_path = os.path.join(case_dir, "input", "connection.txt")
    bond_path = os.path.join(case_dir, "input", "bond.txt")

    if not os.path.exists(coord_path):
        print(f"  [跳过] {case_dir}: 找不到 coord.txt")
        return False

    # 读取 coord.txt 获取网格尺寸
    with open(coord_path, "r", encoding="utf-8") as f:
        lines = f.readlines()
    n_atoms = len(lines) - 1  # 去掉表头
    N = int(n_atoms ** 0.5)
    if N * N != n_atoms:
        print(f"  [跳过] {case_dir}: 非正方形网格 (n_atoms={n_atoms})")
        return False

    print(f"  {case_dir}: {N}x{N} 网格")

    # 读取现有 connection.txt，检查是否已有 k2
    has_k2 = False
    if os.path.exists(conn_path):
        with open(conn_path, "r") as f:
            for line in f:
                if "k2" in line:
                    has_k2 = True
                    break

    if has_k2:
        print(f"    k2 已存在，跳过 connection.txt")
    else:
        # 追加 k2 键到 connection.txt
        with open(conn_path, "a", encoding="utf-8") as f:
            cnt = 0
            for row in range(N):
                for col in range(N):
                    id1 = row * N + col + 1
                    if col + 1 < N and row + 1 < N:
                        f.write(f"{id1} {(row + 1) * N + (col + 1) + 1} k2\n")
                        cnt += 1
                    if col - 1 >= 0 and row + 1 < N:
                        f.write(f"{id1} {(row + 1) * N + (col - 1) + 1} k2\n")
                        cnt += 1
        print(f"    connection.txt: 追加 {cnt} 条 k2 键")

    # 检查 bond.txt 是否有 k2
    has_bond = False
    if os.path.exists(bond_path):
        with open(bond_path, "r") as f:
            for line in f:
                if line.startswith("k2"):
                    has_bond = True
                    break

    if has_bond:
        print(f"    bond.txt: k2 已存在")
    else:
        with open(bond_path, "a", encoding="utf-8") as f:
            f.write("k2 100.0 1.41421356\n")
        print(f"    bond.txt: 追加 k2 定义")

    return True


if __name__ == "__main__":
    targets = []
    if "--all" in sys.argv:
        base = os.path.dirname(os.path.abspath(__file__))
        for d in sorted(os.listdir(base)):
            if d.startswith("case") and os.path.isdir(os.path.join(base, d)):
                targets.append(os.path.join(base, d))
    else:
        for arg in sys.argv[1:]:
            if arg.startswith("--"):
                continue
            p = arg if os.path.isabs(arg) else os.path.join(os.path.dirname(os.path.abspath(__file__)), arg)
            targets.append(p)

    for t in targets:
        add_k2(t)