#!/usr/bin/env python3
"""
Generate a VASP KPOINTS file for band-structure calculations from a CIF file,
using pymatgen’s HighSymmKpath utilities.
"""

import argparse
import sys

from pymatgen.io.cif import CifParser
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from pymatgen.symmetry.bandstructure import HighSymmKpath
from pymatgen.io.vasp.inputs import Kpoints
from fractions import Fraction


def snap(coords, max_den=2):
    # max_den=2 なら 1/2 までの分母で丸め
    return [float(Fraction(c).limit_denominator(max_den)) for c in coords]

def initialize():
    parser = argparse.ArgumentParser(
        description="Generate a VASP KPOINTS for band structure from a CIF."
    )
    parser.add_argument(
        "ciffile",
        help="Path to the input CIF file."
    )
    parser.add_argument(
        "--primitive_on_read",
        action="store_true",
        help="Convert to primitive cell upon reading from CIF file"
    )
    parser.add_argument(
        "-c", "--conventional",
        action="store_true",
        help="Keep cell red from CIF file"
    )
# https://pymatgen.org/pymatgen.symmetry.html より、このフラグはFalseであるべき
    parser.add_argument(
        "--international-monoclinic",
        action="store_true",
        help="Use the international monoclinic convention when standardizing structure."
    )
    parser.add_argument(
        "--keep-site-properties",
        action="store_true",
        help="Keep original site properties when reducing to primitive standard structure."
    )
    parser.add_argument(
        "-f", "--frac_tolerance",
        type=float,
        default=1e-4,
        help="Fractional tolerance in structure parameters (default: 1e-4)."
    )
    parser.add_argument(
        "-d", "--max_den",
        type=int,
        default=4,
        help="Number of digits to round fractional coordinates in reciprocal space (default: 4)."
    )
    parser.add_argument(
        "-s", "--symprec",
        type=float,
        default=1e-4,
        help="Positional tolerance for symmetry detection (default: 1e-4)."
    )
    parser.add_argument(
        "-a", "--angle-tol",
        type=float,
        dest="angle_tolerance",
        default=0.1,
        help="Angular tolerance (degrees) for symmetry detection (default: 0.1)."
    )
    parser.add_argument(
        "--atol",
        type=float,
        default=0.1,
        help="Absolute tolerance used to determine symmetric equivalence of points and lines on the BZ."
    )
    parser.add_argument(
        "-n", "--ndivisions",
        type=int,
        default=20,
        help="Number of k-point divisions along each path segment (default: 20)."
    )
    parser.add_argument(
        "-p", "--path-type",
        choices=["setyawan_curtarolo", "latimer_munro", "hinuma", "all"],
        default="setyawan_curtarolo",
        help="Which high‐symmetry path convention to use (default: setyawan_curtarolo)."
    )
    parser.add_argument(
        "-o", "--output",
        default="KPOINTS",
        help="Output filename for the KPOINTS file (default: KPOINTS)."
    )
    
    return parser

def update_vars(parser):
    return parser.parse_args()

def read_cif(ciffile, frac_tolerance, primitive = True):
    print()
    print(f"Read cif from {ciffile}")
    parser = CifParser(ciffile, frac_tolerance = frac_tolerance)
    structures = parser.parse_structures(primitive = primitive)
    if not structures:
        print(f"Error: no structures found in '{args.ciffile}'")
        return False

    structure = structures[0]
    print(structure)
    analyzer = SpacegroupAnalyzer(structure)
#    spacegroup = analyzer.get_space_group_symbol()
    spg_name, spg_num = structure.get_space_group_info() 
    print(f"Space group: {spg_name} #{spg_num}")

    return structure

def convert_primitive(structure_original, symprec, angle_tolerance, 
                      international_monoclinic, keep_site_properties):
    print()
    print("Standardize and reduce to primitive cell")
    analyzer = SpacegroupAnalyzer(
        structure_original,
        symprec = symprec,
        angle_tolerance = angle_tolerance
    )
    structure = analyzer.get_primitive_standard_structure(
        international_monoclinic = international_monoclinic,
        keep_site_properties = keep_site_properties
    )
    print(structure)
    analyzer = SpacegroupAnalyzer(structure)
#    spacegroup = analyzer.get_space_group_symbol()
    spg_name, spg_num = structure.get_space_group_info() 
    print(f"Space group: {spg_name} #{spg_num}")

    return structure

def generate_kpoints(structure, args):
    print()
    print("Generate k path")
    # 3) Generate high‐symmetry k‐path
    kpath = HighSymmKpath(structure, path_type=args.path_type, 
                symprec = args.symprec, angle_tolerance = args.angle_tolerance, atol = args.atol)
    print("kpath:", kpath.kpath["path"])
    print("K-point path segments:")
    for segment in kpath.kpath["path"]:
        print("  ", " -> ".join(segment))
    print("\nHigh-symmetry k-points:")
    for label, coord in kpath.kpath["kpoints"].items():
        print(f"  {label:10} {coord}")

    print()
    print("Correct k path")
    for name, frac in kpath.kpath['kpoints'].items():
        kpath.kpath['kpoints'][name] = snap(frac, max_den = args.max_den)
    print("High-symmetry k-points:")
    for label, coord in kpath.kpath["kpoints"].items():
        print(f"  {label:10} {coord}")

    print()
    print("Create KPOINTS")
    # 4) Build and write the KPOINTS file in line-mode
    kpoints = Kpoints.automatic_linemode(divisions = args.ndivisions, ibz=kpath)
    print(kpoints)
    kpoints.write_file(args.output)
    print(f"\nWrote KPOINTS to '{args.output}'")


def main():
    parser = initialize()
    args = update_vars(parser)
    #argsの内容を確認
    print("Arguments:")
    for key, value in vars(args).items():
        print(f"  {key}: {value}")

    print()
    print(f"Read CIF file and generate KPOINTS for VASP")
    print(f"CIF file: {args.ciffile}")
    print(f"Convert to primitive cell on read_cif(): {args.primitive_on_read}")
    print(f"Use conventional cell: {args.conventional}")

    structure = read_cif(args.ciffile, args.frac_tolerance, primitive = args.primitive_on_read)
    if structure is None: return None

    if not args.conventional:
        structure = convert_primitive(structure, args.symprec, args.angle_tolerance, 
                      args.international_monoclinic, args.keep_site_properties)
        if structure is None: return None

    generate_kpoints(structure, args)


if __name__ == "__main__":
    main()