"""
pymatgenとtklibを用いてCIFファイルを対称化し、その結果を新しいCIFファイルに保存するスクリプト。

概要:
    入力されたCIFファイルを読み込み、pymatgenのSpacegroupAnalyzerを使用して結晶構造を対称化します。
    対称化された構造から得られる空間群情報、格子定数、原子サイトの情報を抽出し、
    tklibのtkFileモジュールを使用して新しいCIFファイルとして出力します。
    このスクリプトは、pymatgenの強力な対称性解析機能を活用しつつ、tkCrytalのデータ構造と連携して
    出力を生成するハイブリッドなアプローチを採用しています。

詳細説明:
    1. コマンドライン引数から入力CIFファイルを指定します（デフォルトは 'SrTiO3.cif'）。
    2. ログファイルを生成し、標準出力とログファイルへ出力をリダイレクトします。
    3. tkCIFを使って入力CIFファイルを読み込み、基本的な結晶情報を取得します。
    4. pymatgenのStructure.from_fileメソッドで入力CIFファイルを読み込み、Structureオブジェクトを生成します。
    5. SpacegroupAnalyzerを用いてStructureオブジェクトを解析し、対称化されたStructureオブジェクトを取得します。
    6. 対称化された構造から、空間群名、番号、対称操作、格子パラメータ、等価サイトの情報を抽出します。
    7. 抽出した情報（格子、空間群、対称操作、原子サイトの分数座標と占有率）を基に、
       新しいCIFファイル（'{filebody}-symmetrized.cif'）を生成・保存します。
    8. 対称操作のマトリクス成分をCIF形式の文字列（例: 'x', 'x+y', '1/2-z'）に変換するヘルパー関数を使用します。

関連リンク:
    :doc:`symmetrize_tklib_usage`
"""
import re
import sys

from pymatgen.io.cif import CifParser
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
from pymatgen.core.structure import Structure
from pymatgen.core.lattice import Lattice
from pymatgen.core.periodic_table import Element


from tklib.tkapplication import tkApplication
from tklib.tkutils import terminate, pint, pfloat, getarg, getintarg, getfloatarg
from tklib.tkfile import tkFile
from tklib.tkcrystal.tkcif import tkCIF


infile = 'SrTiO3.cif'
prec = 1.0e-3

app    = tkApplication()
argv = sys.argv
narg = len(argv)
if narg >= 2:
    infile = argv[1]


#==========================================
# Main prgram
#==========================================
def main():
    """
    メインプログラムのエントリーポイント。CIFファイルを読み込み、対称化し、結果を出力する。

    詳細説明:
        - 入力CIFファイルを読み込み、`tkCIF`と`pymatgen.Structure`の両方で構造データを取得します。
        - `pymatgen.SpacegroupAnalyzer`を使用して構造を対称化し、空間群情報、格子パラメータ、等価サイトを取得します。
        - 対称化された構造情報（格子、空間群、原子サイト）を基に、新しいCIFファイルを生成して保存します。
        - ログファイルへのリダイレクトやエラーハンドリングも行います。

    :returns: None
    """
    logfile = app.replace_path(infile, template = ["{dirname}", "{filebody}-out.txt"])
    print(f"Open logfile [{logfile}]")
    app.redirect(targets = ["stdout", logfile], mode = 'w')

    convCIFfile = app.replace_path(infile, template = ["{dirname}", "{filebody}-symmetrized.cif"])

    print("")
    print(f"input : {infile}")   
    print(f"log file: {logfile}")
    print(f"output symmetrized CIF file: {convCIFfile}")

#=====================================
# Ready by tkProg
#=====================================
    print("")
    print(f"Read [{infile}] for tkCIF")
    cif = tkCIF()
    cifdata = cif.ReadCIF(infile, find_valid_structure = True)
    if cifdata is None:
        app.terminate(f"Error: Can not read [{infile}]", pause = True)
        exit()
        
    cifdata.Print()
    cry_source = cifdata.GetCrystal()
#    cry_source.print_inf()

#=====================================
# Ready by pymatgen
#=====================================
    print("")
    print(f"Read [{infile}] for pymatgen")
    structure = Structure.from_file(infile)
    print("Structure:", structure)

    print("")
    print(f"Symmetrized:")
    analyzer = SpacegroupAnalyzer(structure)
    symmetrized_structure      = analyzer.get_symmetrized_structure()
    symmetrized_structure_dict = symmetrized_structure.as_dict()

    spg_name, spg_num = symmetrized_structure.get_space_group_info()
    symmetry_ops    = analyzer.get_symmetry_operations()
    symmetry_matrix = [op.as_dict()['matrix'] for op in symmetry_ops]
    nsym = len(symmetry_matrix)
    print("Space group: ", spg_name, spg_num)
    print(f"Number of symmetry options: {nsym}")
#    print("Symmetry operatoins: ", symmetry_matrices)

    print("")
    print("Lattice: ")
    lattice_inf = symmetrized_structure.lattice
    a, b, c, alpha, beta, gamma = lattice_inf.parameters
    volume = lattice_inf.volume
    aij    = lattice_inf.matrix
    print("  Lattice parameters:", a, b, c, alpha, beta, gamma)
    print("  Matrix: ", aij)
    print("  Volume: ", volume)

    print("")
    print("Equivalent sites:")
    eq_sites = symmetrized_structure.equivalent_sites
    for sites in eq_sites:
        composition  = sites[0].species
        for atom_name in composition.keys():
            print("  ", atom_name, sites[0].frac_coords, "  occ=", composition[atom_name])
#        print("dir(sites[0])=", dir(sites[0]))
#        for site in sites:
#            print("  ", site)
    
    print("")
    print("All sites:")
    sites_inf = symmetrized_structure.sites
    for site in sites_inf:
        composition  = site.species
        for atom_name in composition.keys():
            print("  ", atom_name, site.frac_coords, "  occ=", composition[atom_name])

#=====================================
# Save
#=====================================
    print("")
    print(f"Save the symmetrized structure to [{convCIFfile}]")
# This does not save the symmetrized structure. So implement with tkCrytal
#    symmetrized_structure.to(filename = convCIFfile, fmt = "cif")
    out = tkFile(convCIFfile, 'w')
    if out.fp is None:
        app.terminate(f"Error: Can not write to [{convCIFfile}]", pause = True)

    out.write(f"data_{cry_source.path}\n")
    out.write(f"_cell_length_a     {a}\n")
    out.write(f"_cell_length_b     {b}\n")
    out.write(f"_cell_length_c     {c}\n")
    out.write(f"_cell_angle_alpha  {alpha}\n")
    out.write(f"_cell_angle_beta   {beta}\n")
    out.write(f"_cell_angle_gamma  {gamma}\n")
    out.write(f"_cell_volume       {volume}\n")
    out.write(f"\n")
    out.write(f"_symmetry_space_group_name_H-M  '{spg_name}'\n")
    out.write(f"_symmetry_Int_Tables_number     {spg_num}\n")
    out.write(f"_chemical_formula_structural  {cry_source.ChemicalFormula}\n")
    out.write(f"_chemical_formula_sum         {cry_source.ChemicalFormula}\n")
    out.write(f"_cell_formula_units_Z   {cry_source.ChemicalFormulaUnit}\n")

    def to_str(v):
        """
        浮動小数点数をCIF形式の分数の文字列に変換する。

        詳細説明:
            - 特定の浮動小数点値（0, ±1, ±0.5, ±0.25, ±1/3, ±2/3, ±1/6, ±5/6など）を、それぞれの分数表現（例: '1/2', '-1/3'）に変換します。
            - 正の整数または単純な分数ではない場合は、入力値を文字列として返します。
            - 0.0の場合はNoneを返します。

        :param v: float
            変換する浮動小数点数。
        :returns: Optional[str]
            変換された文字列、または変換できない場合は元の値、または0.0の場合はNone。
        """
        eps = 1.0e-6
        if v == 0.0:
            return None
        if v == 1.0:
            return '+'
        if v == -1.0:
            return '-'
        if abs(v - 0.5) < eps:
            return '1/2'
        if abs(v + 0.5) < eps:
            return '-1/2'
        if abs(v - 0.25) < eps:
            return '1/4'
        if abs(v + 0.25) < eps:
            return '-1/4'
        if abs(v - 1.0/3.0) < eps:
            return '1/3'
        if abs(v + 1.0/3.0) < eps:
            return '-1/3'
        if abs(v - 2.0/3.0) < eps:
            return '2/3'
        if abs(v + 2.0/3.0) < eps:
            return '-2/3'
        if abs(v - 1.0/6.0) < eps:
            return '1/6'
        if abs(v + 1.0/6.0) < eps:
            return '-1/6'
        if abs(v - 5.0/6.0) < eps:
            return '5/6'
        if abs(v + 5.0/6.0) < eps:
            return '-5/6'
        if v > 0.0:
            return re.sub(r'^\+', '', v)

        return v
            
    def vector_to_str(v):
        """
        3DベクトルをCIF形式の対称操作文字列（例: 'x', 'x+y', '-x+z'）に変換する。

        詳細説明:
            - 入力された3要素ベクトル（[x, y, z]）の各要素を`to_str`関数で変換し、'x', 'y', 'z'を連結して対称操作の文字列を生成します。
            - 例えば、[1, 0, 0] は 'x' に、[0, 1, 0] は 'y' に、[0, 0, -1] は '-z' に、[1, 1, 0] は 'x+y' になります。
            - 先頭の'+'は除去されます。

        :param v: List[float]
            3要素の浮動小数点数ベクトル。
        :returns: str
            対称操作を表す文字列。
        """
        sx = to_str(v[0])
        sy = to_str(v[1])
        sz = to_str(v[2])
        s = ''
        if sx:
            s += sx + 'x'
        if sy:
            s += sy + 'y'
        if sz:
            s += sz + 'z'

        return re.sub(r'^\+', '', s)

    def matrix_to_str(m):
        """
        3x3行列をCIF形式の対称操作文字列のタプル（'x', 'y', 'z'成分）に変換する。

        詳細説明:
            - 入力された3x3行列の各行（ベクトル）を`vector_to_str`関数に渡し、それぞれの変換結果をタプルとして返します。
            - これは、CIFの `_symmetry_equiv_pos_as_xyz` エントリの 'x, y, z' 部分を生成するために使用されます。

        :param m: List[List[float]]
            3x3の浮動小数点数行列。
        :returns: Tuple[str, str, str]
            'x', 'y', 'z'成分を表す文字列のタプル。
        """
        x = vector_to_str(m[0])
        y = vector_to_str(m[1])
        z = vector_to_str(m[2])

        return x, y, z

    out.write(f"\n")
    out.write(f"loop_\n")
    out.write(f" _symmetry_equiv_pos_site_id\n")
    out.write(f" _symmetry_equiv_pos_as_xyz\n")
    for i in range(nsym):
        x, y, z = matrix_to_str(symmetry_matrix[i])
        out.write(f" {i+1:2}  '{x}, {y}, {z}'\n")

    out.write(f"\n")
    out.write(f"loop_\n")
    out.write(f" _atom_site_type_symbol\n")
#    out.write(f" _atom_site_label\n")
    out.write(f" _atom_site_fract_x\n")
    out.write(f" _atom_site_fract_y\n")
    out.write(f" _atom_site_fract_z\n")
    out.write(f" _atom_site_occupancy\n")
    for sites in eq_sites:
        composition  = sites[0].species
        for atom_name in composition.keys():
            pos = sites[0].frac_coords
            occ = composition[atom_name]
            
            out.write(f" {atom_name}")
            out.write(f"    {pos[0]:12.8f}  {pos[1]:12.8f}  {pos[2]:12.8f}")
            out.write(f"  {occ}\n")

    out.close()


    app.terminate(pause = True)


if __name__ == "__main__":
    main()