"""
CIFファイルから結晶構造情報を読み込み、pymatgenを用いて解析・表示するスクリプト。

概要:
    本スクリプトは、指定されたCIFファイルから結晶構造データを読み込み、
    pymatgenライブラリの機能を用いて、その構造に関する詳細な情報を解析し、標準出力に表示します。
    具体的には、空間群の情報、格子定数、対称操作、原子サイトの座標と組成などが含まれます。
    また、元の構造に加え、対称化された構造の解析結果も提示します。

詳細説明:
    1.  コマンドライン引数またはデフォルト設定に基づいて入力CIFファイルパスを決定します。
    2.  `tklib.tkApplication` を使用して、標準出力に加え、ログファイルへの出力リダイレクトを設定します。
    3.  pymatgenの`Structure.from_file`メソッドでCIFファイルを読み込み、`Structure`オブジェクトを作成します。
    4.  作成された`Structure`オブジェクトについて、`print_inf`関数を呼び出して詳細情報を表示します。
    5.  `SpacegroupAnalyzer`を用いて構造を対称化し、その対称化された`Structure`オブジェクトについても
        同様に`print_inf`関数で詳細情報を表示します。
    6.  表示される情報には、空間群のシンボルと番号、対称操作の行列、格子定数、格子ベクトル、
        単位胞体積、独立なサイトおよび全サイトの原子種、分数座標、占有率が含まれます。

関連リンク:
    :doc:`cif_inf_pymatgen_usage`
"""
from pprint import pprint
import re

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


infile = 'SrTiO3.cif'


app    = tkApplication()
infile      = getarg( 1, infile)


#==========================================
# Main prgram
#==========================================
def print_matrix(m):
    """
    3x3行列を整形して標準出力に表示します。

    概要:
        与えられた3x3の数値行列を行と列にフォーマットして表示します。

    詳細説明:
        行列の各要素は小数点以下4桁の浮動小数点数としてフォーマットされ、
        パイプ記号で囲まれた見やすい形式で出力されます。

    :param m: 表示する3x3の数値行列。リストのリスト形式で指定します。
    :type m: list[list[float]]
    :returns: None
    """
    print(f"| {m[0][0]:8.4f} {m[0][1]:8.4f} {m[0][2]:8.4f} |")
    print(f"| {m[1][0]:8.4f} {m[1][1]:8.4f} {m[1][2]:8.4f} |")
    print(f"| {m[2][0]:8.4f} {m[2][1]:8.4f} {m[2][2]:8.4f} |")

def print_inf(structure):
    """
    pymatgenのStructureオブジェクトから詳細な結晶構造情報を抽出し、表示します。

    概要:
        与えられたStructureオブジェクトの空間群、対称操作、格子情報、およびサイト情報を整形して表示します。

    詳細説明:
        以下の情報が表示されます。
        - 構造全体の概要 (`structure.__str__`)
        - 空間群のシンボルと番号
        - 全ての対称操作の数と各対称操作の行列
        - 格子定数 (a, b, c, alpha, beta, gamma)
        - 格子ベクトル
        - 単位胞の体積
        - 独立なサイトの原子種と分数座標
        - 全てのサイトの原子種、分数座標、および占有率

    :param structure: 解析対象のpymatgen Structureオブジェクト。
    :type structure: pymatgen.core.structure.Structure
    :returns: None
    """
    print(structure)

    structure_dict = structure.as_dict()
    print("structure keys:", structure_dict.keys())
    structure_inf  = structure_dict.get('structure', None)
    spacegroup_inf = structure_dict.get('spacegroup', None)
    charge_inf     = structure_dict.get('charge', None)
#    print("  structure: ", structure_inf)
#    print("    keys:", structure_inf.keys())
#    print("  spacegroup: ", spacegroup_inf)
#    print("  charge : ", charge_inf)

    analyzer = SpacegroupAnalyzer(structure)
#    spacegroup = analyzer.get_space_group_symbol()
    spg_name, spg_num = 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("")
    print(f"Space group: {spg_name} #{spg_num}")
    print(f"Symmetry operatoins: {nsym}")
    for i in range(nsym):
        print(f"op #{i+1}")
        print_matrix(symmetry_matrix[i])

    print("")
    print("Lattice: ")
    lattice_inf = 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("  Lattice vectors:")
    pprint(aij)
    print("  Volume: ", volume)
#    lattice_inf = structure_dict['structure']['lattice']
#    a, b, c            = structure.lattice.abc
#    alpha, beta, gamma = structure.lattice.angles

    """
    print("")
    print("Species  : ")
    species_inf = structure.species
    print("dir(spec)=", dir(species_inf))
#    print("vars(spec)=", vars(species_inf))
    print("species_inf=", species_inf)
#    print("dir(species_inf)=", dir(species_inf))
#    print("vars(species_inf)=", vars(species_inf))
    mat_spec_occ = symmetrized_structure.species_and_occu
    print("mat_spec_occ:")
    print("mat_spec_occ=", mat_spec_occ)
#    for occ in mat_spec_occ:
#        print("  occ:", occ.to_data_dict)
#        print("  occ:", dir(occ))
    exit()
    """

    print("")
    eq_sites = getattr(structure, "equivalent_sites", None)
    if eq_sites is None:
        print("No information for equivalent sites")
    else:
        print("Independent sites:")
        for sites in eq_sites:
            print("  ", sites[0].species, sites[0].frac_coords)
# 以降のデータは等価位置なので表示しない        
#        print("dir(sites[0])=", dir(sites[0]))
#        for site in sites:
#            print("  ", site)

    print("")
    print("All sites:")
    sites_inf = structure.sites
    for site in sites_inf:
#        print("  ", site)
#    print(sites_inf[0].coords)
        composition = site.species
        for atom_name in composition.keys():
            print("  ", atom_name, site.frac_coords, "  occ=", composition[atom_name])


def main():
    """
    プログラムの主要な実行フローを定義します。

    概要:
        指定されたCIFファイルを読み込み、元の構造と対称化された構造の両方について詳細情報を表示します。

    詳細説明:
        1. 入力CIFファイル名からログファイル名と対称化されたCIFファイルの出力パスを生成します。
        2. 標準出力をログファイルにもリダイレクトします。
        3. 入力CIFファイルを`pymatgen.core.structure.Structure`オブジェクトとして読み込みます。
        4. 読み込んだ元の構造の情報を`print_inf`関数を使って表示します。
        5. `SpacegroupAnalyzer`を使用して構造を対称化し、その対称化された構造の情報を`print_inf`関数を使って表示します。
        6. プログラムの実行終了時に一時停止します。

    :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}")

    print("")
    print(f"Read [{infile}]")
    structure = Structure.from_file(infile)
    
    print("")
    print("Structure as original input:")
    print_inf(structure)

    print("")
    print("Symmetrized structure:")
    analyzer = SpacegroupAnalyzer(structure)
    symmetrized_structure = analyzer.get_symmetrized_structure()
    print_inf(symmetrized_structure)

    app.terminate(pause = True)


if __name__ == "__main__":
    main()