"""
VASPの分子動力学計算出力ファイルを解析し、XCrySDen AXSFファイルへの変換や、平均二乗変位 (MSD) の履歴を計算するスクリプトです。

詳細説明:
    VASPのOUTCARファイルから時系列の結晶構造を読み込み、XCrySDen形式のアニメーションファイル (AXSF) を生成します。
    また、MDシミュレーションの履歴情報（温度、エネルギー、格子定数、原子位置、平均二乗変位 (MSD) など）を抽出し、CSV形式で出力します。

関連リンク:
    :doc:`make_md_history_usage`
"""
import os
import sys
import shutil
import copy
import glob
import csv
import numpy as np
from numpy import exp, log, sin, cos, tan, arcsin, arccos, arctan, pi

sys.path.append("c:/Programs/python/lib")
sys.path.append("d:/Programs/python/lib")

from tklib.tkfile import tkFile
import tklib.tkre
from tklib.tkutils import IsDir, IsFile, SplitFilePath
from tklib.tkutils import terminate, pint, pfloat, getarg, getintarg, getfloatarg, safe_getelement
from tklib.tksci.tksci import Reduce01, Round
from tklib.tksci.tkmatrix import make_matrix1, make_matrix2, make_matrix3
from tklib.tkcrystal.tkcif import tkCIF, tkCIFData
from tklib.tkcrystal.tkatomtype import tkAtomType
from tklib.tkcrystal.tkatomtypeobject import GetAtomName, ReadAtomDB, GetAtomInformation
from tklib.tkcrystal.tkcrystal import tkCrystal
from tklib.tkcrystal.tkvasp import tkVASP


"""
Convert VASP MD output files to XCrySDen AXSF file
"""

#================================
# global parameters
#================================
debug = 0

# input path
CAR_path     = 'OUTCAR'

cifpath = 'test.cif'
averagecsvfile = None
historycsvfile = None

# # of sites whose positions are written in averagecsvfile
nxyzout = 5

# cif read configuration
single = 1
findvalidstructure = 1

# Initial positions: 'POSCAR', 'OUTCAR', 'average'
#origin = 'POSCAR'
#origin = 'OUTCAR'
origin = 'average'

# Drift correction: 'no', 'cm'
drift_correction = 'no'
#drift_correction = 'cm'

# mode: 'site', 'atom'
mode = 'atom'

# # of MD steps to be removed to calculate MSD
nskip = 0

# Maximum # of MD steps to be read
nmax_read = -1

#=============================
# Treat argments
#=============================
def usage():
    """
    概要: スクリプトの正しい使用方法を標準出力に表示します。

    詳細説明:
        コマンドライン引数の種類によって異なる使用例を示します。
        引数不足や不正なモード指定があった場合に呼び出され、ユーザーに適切な操作方法を案内します。

    引数:
        なし

    戻り値:
        なし
    """
    global CAR_path

    print("")
    print("Usage:")
    print("  (a) Create XSF file from OUTCAR")
    print("     python {} CAR_path init nmax_read".format(sys.argv[0]))
    print("     ex: python {} {} {} {} {}".format(sys.argv[0], CAR_path, mode, origin, nskip, nmax_read))
    print("  (b) Calculate MSD")
    print("     python {} CAR_path mode origin drift_correctin nskip nmax_read".format(sys.argv[0]))
    print("         mode  : 'atom', 'site'")
    print("         origin: 'POSCAR', 'OUTCAR', 'average'")
    print("         drift_correctin: 'no', 'cm' (by center of mass)")
    print("     ex: python {} {} {} {} {} {} {}".format(sys.argv[0], CAR_path, mode, origin, drift_correction, nskip, nmax_read))
    
    
def updatevars():
    """
    概要: コマンドライン引数を解析し、グローバル設定変数を更新します。

    詳細説明:
        `sys.argv` からコマンドライン引数を取得し、
        `CAR_path`, `mode`, `origin`, `drift_correction`, `nskip`, `nmax_read`
        などのグローバル変数の値を更新します。
        引数が不足している場合は `usage()` 関数を呼び出してスクリプトを終了します。

    引数:
        なし (コマンドライン引数を内部で処理します)

    戻り値:
        なし
    """
    global CAR_path, mode, origin, drift_correction, nskip, nmax_read

    if getarg(1, None) is None:
        terminate("", usage = usage)
        
    CAR_path  = getarg(1, CAR_path)
    mode      = getarg(2, mode)
    if mode == 'init':
        nmax_read = getintarg(3, nmax_read)
    elif mode == 'atom' or mode == 'site':
        origin          = getarg(3, origin)
        drift_correctin = getarg(4, nskip)
        nskip           = getintarg(5, nskip)
        nmax_read       = getintarg(6, nmax_read)

def WriteXSFAnimationFileHeader(outfile, nStep):
    """
    概要: XCrySDen AXSFアニメーションファイルのヘッダーを書き込みます。

    詳細説明:
        AXSFファイルのアニメーションステップ数 (`ANIMSTEPS`) と
        結晶構造の開始 (`CRYSTAL`) を示すヘッダー情報を指定されたファイルに書き込みます。

    引数:
        :param outfile: str: 出力XSFファイルのパス。
        :param nStep: int: アニメーションの総ステップ数。

    戻り値:
        :returns: int or None: ファイルの書き込みに成功した場合は 1、失敗した場合は None を返します。
    """
    xsf = tkFile(outfile, 'w')
    if not xsf:
        return None

    xsf.Write("ANIMSTEPS {}\n".format(nStep))
    xsf.Write("CRYSTAL\n")

    xsf.Close()
    return 1

def WriteXSFFileFromCrystal(outfile, iCycle, cry):
    """
    概要: 指定された結晶構造情報をXCrySDen AXSFファイルに追記します。

    詳細説明:
        特定のMDステップにおける格子ベクトル (`PRIMVEC`, `CONVVEC`) と
        各原子のカート座標および力または速度 (`PRIMCOORD`) を
        指定されたXSFファイルに追記します。
        原子の情報を基に原子番号を決定します。

    引数:
        :param outfile: str: 追記先のXSFファイルのパス。
        :param iCycle: int: 現在のMDステップ（サイクル）番号。
        :param cry: tkcrystal.tkcrystal.tkCrystal: 書き込む結晶構造オブジェクト。

    戻り値:
        :returns: int or None: ファイルの書き込みに成功した場合は 1、失敗した場合は None を返します。
    """
    types = cry.AtomTypeList()

    xsf = tkFile(outfile, 'a')
    if not xsf:
        return None

    aij = cry.LatticeVectors()
    xsf.Write("PRIMVEC {}\n".format(iCycle))
    xsf.Write(" %16.10f %16.10f %16.10f\n" % (aij[0][0], aij[0][1], aij[0][2]))
    xsf.Write(" %16.10f %16.10f %16.10f\n" % (aij[1][0], aij[1][1], aij[1][2]))
    xsf.Write(" %16.10f %16.10f %16.10f\n" % (aij[2][0], aij[2][1], aij[2][2]))
    xsf.Write("CONVVEC {}\n".format(iCycle))
    xsf.Write(" %16.10f %16.10f %16.10f\n" % (aij[0][0], aij[0][1], aij[0][2]))
    xsf.Write(" %16.10f %16.10f %16.10f\n" % (aij[1][0], aij[1][1], aij[1][2]))
    xsf.Write(" %16.10f %16.10f %16.10f\n" % (aij[2][0], aij[2][1], aij[2][2]))

    types        = cry.AtomTypeList()
    ntypes       = len(types)
    sites        = cry.ExpandedAtomSiteList()
    nsites       = len(sites)

    xsf.Write("PRIMCOORD {}\n".format(iCycle))
    xsf.Write(" %d  1\n" % (nsites))
#$Crystal2 != ''の時、重心の変化を調べる
    centx1, centy1, centz1 = 0.0, 0.0, 0.0
    centx2, centy2, centz2 = 0.0, 0.0, 0.0

    for i in range(nsites):
        atom         = sites[i]
        name         = atom.AtomNameOnly()
        itype        = atom.iAtomType()
        AtomicNumber = types[itype].AtomicNumber()
        if AtomicNumber == 0:
            AtomicNumber = ReadAtomDB(name)
        
        x, y, z    = atom.Position(1)
        fx, fy, fz = atom.Force()
        if fx is None:
            fx, fy, fz = atom.Velocity()

        xc, yc, zc = cry.FractionalToCartesian(x, y, z)
        if fz is None:
            fx = fy = fz = 0.0
        xsf.Write("%3d   %12.7f %12.7f %12.7f   %12.7f %12.7f %12.7f\n" 
                % (AtomicNumber, xc, yc, zc, fx, fy, fz))

    xsf.Close()

    return 1


def read_axsf(xsffile):
    """
    概要: XCrySDen AXSFファイルを読み込み、結晶構造のリストを生成します。

    詳細説明:
        指定されたAXSFファイルから各MDステップの格子情報と原子座標を解析し、
        `tkCrystal` オブジェクトのリストとして返します。
        `nmax_read` (グローバル変数) が設定されている場合、そのステップ数を超えるデータは読み込まれません。

    引数:
        :param xsffile: str: 読み込むAXSFファイルのパス。

    戻り値:
        :returns: list[tkcrystal.tkcrystal.tkCrystal] or None: 読み込まれた`tkCrystal`オブジェクトのリスト、
                                                               またはファイル読み込み失敗時はNone。
    """
    print("")
    print("read_axsf(): Read [{}]".format(xsffile))

    fp = open(xsffile)
    if not fp:
        return None
    
    line = fp.readline().split()
    nsteps = pint(line[1])
    print("nsteps:", nsteps)
#CRYSTAL
    line = fp.readline().split()

    crylist = []
    for istep in range(nsteps):
        if istep % 50 == 0:
            print(" {}".format(istep), flush = True, end = '')

        cry = tkCrystal()
        cry.path = xsffile
#PRIMVEC 0
#     7.9137700000     0.0000000000     0.0000000000
#     0.0000000004     7.9092200000     0.0000000000
#     0.0000000004     0.0000000004     7.9057700000
#CONVVEC 0
        for j in range(5):
            fp.readline()

#     7.9137700000     0.0000000000     0.0000000000
#     0.0000000004     7.9092200000     0.0000000000
#     0.0000000004     0.0000000004     7.9057700000
        aij = make_matrix1(3, [])
        aij[0] = fp.readline().split()
        aij[1] = fp.readline().split()
        aij[2] = fp.readline().split()
        for i in range(3):
            for j in range(3):
                aij[i][j] = pfloat(aij[i][j])
        cry.SetLatticeVectors(aij)
#        latt = cry.LatticeParameters()
#        print("latt:", latt)

#PRIMCOORD 0
        fp.readline()
# 40  1
        nsites, id = fp.readline().split()
        nsites = pint(nsites)
        for isite in range(nsites):
# 38      3.9546100    5.9361100    5.8644600      0.0000850   -0.0003722    0.0007310
            ai = fp.readline().split()
#        inf = self.GetAtomInformation(atomname.capitalize())
            AtomicNumber = pint(ai[0])
            atomname = GetAtomName(AtomicNumber)
            xyz = [pfloat(ai[1]), pfloat(ai[2]), pfloat(ai[3])]
            frac = cry.CartesianToFractional(*xyz)
#            xyz2 = cry.FractionalToCartesian(*frac)
#            print("  {}: #{} {} ({}, {}, {})  ({}, {}, {})".format(isite, AtomicNumber, atomname, *xyz, *frac))
            cry.AddAtomSite(label = None, name = atomname, pos = frac)
        
        cry.ExpandCoordinates()
#        cry.PrintInf()
#        exit()

        crylist.append(cry)

        if nmax_read >= 0 and istep > nmax_read:
            print("")
            print("=========================================================")
            print("*** MD steps exceeds nmax_read ({}).".format(nmax_read))
            print("    Exit loop.")
            print("=========================================================")
            print("")
            break


    return crylist

def make_axsf():
    """
    概要: VASPのOUTCARファイルからXCrySDen AXSFファイルを生成します。

    詳細説明:
        `CAR_path` (グローバル変数) で指定されたディレクトリ内のVASP出力ファイル
        （OUTCAR, INCAR, POSCAR）を読み込みます。
        OUTCARから時系列の結晶構造を抽出し、`md.xsf` というXCrySDen形式のアニメーションファイルと、
        MD情報（温度、エネルギーなど）を記録した `inf.csv` ファイルを生成します。
        `nmax_read` (グローバル変数) で読み込む最大ステップ数を制限できます。

    引数:
        なし (グローバル変数 `CAR_path`, `nskip`, `nmax_read` を使用)

    戻り値:
        :returns: list[tkcrystal.tkcrystal.tkCrystal]: 読み込まれた結晶構造のリスト。
    """
    global CAR_path
    
    vasp = tkVASP()
    
    CAR_path = vasp.getdir(CAR_path)

    print("")
    print("mode:", mode)
    print("# of MD steps to be removed to calculate MSD: {}".format(nskip))
    print("Maximum # of MD steps to be read            : {}".format(nmax_read))

    print("")
    INCAR_path   = vasp.get_INCAR(CAR_path)
    POSCAR_path  = vasp.get_POSCAR(CAR_path)
    CONTCAR_path = vasp.get_CONTCAR(CAR_path)
    OUTCAR_path  = vasp.get_OUTCAR(CAR_path)
    dirname, basename, filebody, ext = SplitFilePath(INCAR_path)
    xsffile = os.path.join(dirname, 'md.xsf')
    inffile = os.path.join(dirname, 'inf.csv')

    print("CAR dir: ", CAR_path)
    print("INCAR  : ", INCAR_path)
    print("POSCAR : ", POSCAR_path)
    print("CONTCAR: ", CONTCAR_path)
    print("OUTCAR : ", OUTCAR_path)
    print("XSF    : ", xsffile)
    print("INF    : ", inffile)
    print("")
    print("Mode: {}".format(mode))

    print("")
    print("*** Read initial crystal structure from [{}]".format(POSCAR_path))
    initial_cry = vasp.read_poscar(POSCAR_path)
    a, b, c, alpha, beta, gamm = initial_cry.LatticeParameters()
    print("cell: {:12.8f} {:12.8f} {:12.8f} A   {:10.6f} {:10.6f} {:10.6f}".format(a, b, c, alpha, beta, gamm))
    initial_cry.PrintInf()

    print("")
    print("Read INCAR [{}]".format(INCAR_path))
    incar_inf = vasp.read_incar_inf(INCAR_path)
    if not incar_inf:
        terminate("Error: Can not read [{}]".format(INCAR_path), usage = usage)

    sample_name = safe_getelement(incar_inf, 'SYSTEM', '') #incar_inf['SYSTEM']
    if not sample_name:
        sample_name = 'hogehoge'
    dt = pfloat(safe_getelement(incar_inf, 'POTIM', 0.0))  #incar_inf['POTIM'])
    print("")
    print("sample_name:", sample_name)
    print("POTIM: {} fs".format(dt))

    print("")
    print("Read crystal structures from [{}]".format(OUTCAR_path))
    crylist = [initial_cry]
    inflist = []

    fp = tkFile(OUTCAR_path, 'r')
    if not fp:
        terminate("Error: Can not read [{}]".format(outcarpath), usage = usage)
    i = 0
    print("  Read structures ", end = '')
    while 1:
        inf, cry = vasp.read_next_crystalstructure(fp, initial_cry, IsReduce01 = False)
        if not cry:
            break

#        cry.PrintInf()
        if i % 50 == 0:
            print(" {}".format(i), flush = True, end = '')
        crylist.append(cry)
        if i == 0:
            inflist.append(inf)
        inflist.append(inf)
        i += 1

        if nmax_read >= 0 and i > nmax_read:
            print("")
            print("=========================================================")
            print("*** MD steps exceeds nmax_read ({}).".format(nmax_read))
            print("    Exit loop.")
            print("=========================================================")
            print("")
            break

    print("")
    fp.Close()
    
    inffp = tkFile(inffile, 'w')
    inffp.Write("temperature,TOTEN,EKIN,EKIN_LAT,ETOTAL\n")
    for istep in range(len(inflist)):
        temperature = safe_getelement(inf, "temperature", 0.0)
        TOTEN    = safe_getelement(inflist[istep], "TOTEN", 0.0)
        EKIN     = safe_getelement(inflist[istep], "EKIN", 0.0)
        EKIN_LAT = safe_getelement(inflist[istep], "EKIN_LAT", 0.0)
        ETOTAL   = safe_getelement(inflist[istep], "ETOTAL", 0.0)
        inffp.Write("{},{},{},{},{}\n".format(temperature, TOTEN, EKIN, EKIN_LAT, ETOTAL))
    inffp.Close()
    
    nstructures = len(crylist)
    print("# of crystal structures:", nstructures)

    types        = initial_cry.AtomTypeList()
    ntypes       = len(types)
    sites        = initial_cry.ExpandedAtomSiteList()
    nsites       = len(sites)
    poslist0     = initial_cry.get_position_list(mode = 'all', IsReduce01 = True)
    atomnamelist = initial_cry.get_atom_name_list(mode = 'all', NameOnly = True)
    offset       = make_matrix2(nsites, 3, defval = 0.0)

    nsitesfortypes = make_matrix1(ntypes, defval = 0)
    for i in range(nsites):
        itype = sites[i].iAtomType()
        nsitesfortypes[itype] += 1
#    print("Atom types:", atomnamelist)
    print("# of atom sites for atom types = ", nsitesfortypes)

    if not WriteXSFAnimationFileHeader(xsffile, nstructures):
        terminate("", usage = usage)

    print("  Write structures to [{}]".format(xsffile), end = '')
    for i in range(nstructures):
        if i % 50 == 0:
            print(" {}".format(i), flush = True, end = '')
        cry = crylist[i]
        if not cry:
            break

        WriteXSFFileFromCrystal(xsffile, i, cry)
    print("")

    return crylist


def make_history():
    """
    概要: VASPのMD計算結果から、原子の変位履歴や平均二乗変位 (MSD) を計算し、CSVファイルに出力します。

    詳細説明:
        `CAR_path` (グローバル変数) で指定されたVASP出力ディレクトリから、
        初期・最終構造のCIFファイル、原子ごとの平均位置、
        およびMD履歴データ（格子情報、エネルギー、MSD）を生成します。
        MSDの計算は、`origin` (初期位置の基準), `drift_correction` (ドリフト補正方法),
        `mode` (MSD計算の対象) のグローバル設定に従って行われます。
        必要に応じてAXSFファイルを再生成します。
        計算結果は `*-average.csv` および `*-history.csv` として出力されます。

    引数:
        なし (グローバル変数 `CAR_path`, `mode`, `origin`, `drift_correction`, `nskip`, `nmax_read` を使用)

    戻り値:
        なし (計算結果をファイルに出力し、スクリプトを終了します)
    """
    global CAR_path
    
    vasp = tkVASP()
    
    CAR_path = vasp.getdir(CAR_path)

    print("")
    INCAR_path   = vasp.get_INCAR(CAR_path)
    POSCAR_path  = vasp.get_POSCAR(CAR_path)
    CONTCAR_path = vasp.get_CONTCAR(CAR_path)
    OUTCAR_path  = vasp.get_OUTCAR(CAR_path)
    dirname, basename, filebody, ext = SplitFilePath(INCAR_path)
    xsffile = os.path.join(dirname, 'md.xsf')
    inffile = os.path.join(dirname, 'inf.csv')
    print("CAR dir: ", CAR_path)
    print("INCAR  : ", INCAR_path)
    print("POSCAR : ", POSCAR_path)
    print("CONTCAR: ", CONTCAR_path)
    print("OUTCAR : ", OUTCAR_path)
    print("XSF    : ", xsffile)
    print("INF    : ", inffile)

    print("")
    print("Mode  : {}".format(mode))
    print("Origin: {}".format(origin))
    print("Drift correction: {}".format(drift_correction))
    print("# of MD steps to be removed to calculate MSD: {}".format(nskip))
    print("Maximum # of MD steps to be read            : {}".format(nmax_read))

    print("")
    print("Read INCAR from [{}]".format(INCAR_path))
    incar_inf = vasp.read_incar_inf(INCAR_path)
    if not incar_inf:
        terminate("Error: Can not read [{}]".format(INCAR_path), usage = usage)

    sample_name = incar_inf['SYSTEM']
    if not sample_name:
        sample_name = 'hogehoge'
    dt = pfloat(incar_inf['POTIM'])

    print("")
    print("sample_name:", sample_name)
    print("POTIM: {} fs".format(dt))

    initial_cifpath = sample_name + '-initial.cif'
    final_cifpath   = sample_name + '-final.cif'
    averagecsvfile  = sample_name + '-average.csv'
    historycsvfile  = sample_name + '-history.csv'
    
    print("")
    print("*** Read initial crystal structure from [{}]".format(POSCAR_path))
    initial_cry = vasp.read_poscar(POSCAR_path)
#    initial_cry.PrintInf("cell")
    a, b, c, alpha, beta, gamm = initial_cry.LatticeParameters()
    print("cell: {:12.8f} {:12.8f} {:12.8f} A   {:10.6f} {:10.6f} {:10.6f}".format(a, b, c, alpha, beta, gamm))

    cif = tkCIFData()
    print("Write initial CIF to [{}]".format(initial_cifpath))
    cif.CreateCIFFileFromCCrystal(initial_cry, initial_cifpath)

    print("")
    print("*** Read final crystal structure from [{}]".format(CONTCAR_path))
    final_cry = vasp.read_poscar(CONTCAR_path)
#    final_cry.PrintInf("cell")
    a, b, c, alpha, beta, gamm = final_cry.LatticeParameters()
    print("cell: {:12.8f} {:12.8f} {:12.8f} A   {:10.6f} {:10.6f} {:10.6f}".format(a, b, c, alpha, beta, gamm))

    print("Write final CIF to [{}]".format(final_cifpath))
    cif = tkCIFData()
    cif.CreateCIFFileFromCCrystal(final_cry, final_cifpath)

    print("")
    print("Read trajectory from XSF file [{}]".format(xsffile))
    if not IsFile(xsffile):
        print("")
        print("*** XSF file [{}] does not exist.".format(xsffile))
        print("    Create it from OUTCAR [{}]".format(OUTCAR_path))
        crylist = make_axsf()
    else:
        crylist = read_axsf(xsffile)
    nstructures = len(crylist)
    print("nstep:", nstructures)

    types        = initial_cry.AtomTypeList()
    ntypes       = len(types)
    sites        = initial_cry.ExpandedAtomSiteList()
    nsites       = len(sites)
    atomnamelist = initial_cry.get_atom_name_list(mode = 'all', NameOnly = True)

# Mass related values
    massdic = {}
    for itype in range(ntypes):
        type = types[itype]
        atomname = type.AtomNameOnly()
        inf = GetAtomInformation(atomname)
        massdic[atomname] = inf['mass']
    Mtot = 0.0
    for isite in range(nsites):
        site     = sites[isite]
        atomname = site.AtomNameOnly()
        Mtot += massdic[atomname]
    print("Mtot:", Mtot)

    ntypesforsites = make_matrix1(ntypes, defval = 0)
    for i in range(nsites):
        itype = sites[i].iAtomType()
        ntypesforsites[itype] += 1
    print("# of atomtypes for sites = ", ntypesforsites)

#=========================================================================
# Calculate center of mass and average r and write to average file
#=========================================================================
    print("")
    print("Calculate center of mass and average r")

# List variable of center of mass
    rglist      = []
    poslist0    = initial_cry.get_position_list(mode = 'all', IsReduce01 = False)
    prevposlist = initial_cry.get_position_list(mode = 'all', IsReduce01 = False)
    offset      = make_matrix3(nstructures, nsites, 3, defval = 0.0)
    for istep in range(nstructures):
        t = dt * istep

        cry   = crylist[istep]
#        cry.PrintInf()
        sites   = cry.ExpandedAtomSiteList()
        poslist = cry.get_position_list(mode = 'all', IsReduce01 = False)

# Center of mass
        rg = make_matrix1(3, 0.0)
        for isite in range(nsites):
            site     = sites[isite]
            atomname = site.AtomNameOnly()
            M        = massdic[atomname]
            pos      = site.Position()

# wrapされた座標から、実際の変位量を計算する
# ただし、一単位格子以上の変位があると、変位量は不明になる。msdの計算はあきらめる
            r, xoff, yoff, zoff = cry.GetNearestInterAtomicDistanceWithOffset(prevposlist[isite], pos, AllowZero = True)
#            print(" site#", isite, ": off:", xoff, yoff, zoff, "  offset:", *offset[istep][isite], end = '')
            if istep == 0:
                offset[istep][isite][0] = xoff
                offset[istep][isite][1] = yoff
                offset[istep][isite][2] = zoff
            else:
                offset[istep][isite][0] = offset[istep-1][isite][0] + xoff
                offset[istep][isite][1] = offset[istep-1][isite][1] + yoff
                offset[istep][isite][2] = offset[istep-1][isite][2] + zoff
#            print(" => ", *offset[istep][isite])

            for i in range(3):
                x = pos[i] - offset[istep][isite][i]
                rg[i] += M * x

        for i in range(3):
            rg[i] /= nsites * Mtot
        rglist.append(rg)

        prevposlist = poslist

#==========================================================
# Write to average file
#==========================================================
    print("Write debug-purpose average file to [{}]".format(averagecsvfile))
    fpavg = tkFile(averagecsvfile, 'w')
    if not fpavg:
        terminate("Error: Can not write to [{}]".format(averagecsvfile))
    fpavg.Write("step,t(fs),xg,yg,zg,xoff(0),yoff(0),zoff(0),xoff(1),yoff(1),zoff(1)")
    for i in range(nxyzout):
        fpavg.Write(",<x({})>,<y({})>,<z({})>".format(i, i, i))
    fpavg.Write("\n")

# Average position
    rtot     = make_matrix2(nsites, 3, 0.0)
    ravg     = make_matrix2(nsites, 3, None)
    for istep in range(nstructures):
        t = dt * istep
        fpavg.Write("{},{}".format(istep, t))

        cry   = crylist[istep]
#        cry.PrintInf()
        sites   = cry.ExpandedAtomSiteList()
        poslist = cry.get_position_list(mode = 'all', IsReduce01 = False)

        for isite in range(nsites):
            site     = sites[isite]
            atomname = site.AtomNameOnly()
            pos      = site.Position()

            for i in range(3):
                if drift_correction == 'no':
                    x = pos[i] - offset[istep][isite][i]
                elif drift_correction == 'cm':
                    x = pos[i] - offset[istep][isite][i] - rglist[istep][i]
                else:
                    terminate("Error: Invalid drift_correction [{}]".format(drift_correction))
                
                if istep <= nskip:
                    ravg[isite][i] = x
                else:
                    rtot[isite][i] += x
                    ravg[isite][i] = rtot[isite][i] / (istep - nskip)

#        print("  {:04d}: x=({}, {}, {}) rg=({}, {}, {})".format(istep, *pos, *rg))
        print("rg:",rglist[istep][0], rglist[istep][1], rglist[istep][2])
        fpavg.Write(",{},{},{}".format(rglist[istep][0], rglist[istep][1], rglist[istep][2]))
        fpavg.Write(",{},{},{},{},{},{}".format(*offset[istep][0], *offset[istep][1]))

        for isite in range(nxyzout):
            fpavg.Write(",{},{},{}".format(ravg[isite][0], ravg[isite][1], ravg[isite][2]))
        fpavg.Write("\n")

    fpavg.Close()

    print("Offset shift (center of mass):")
    print("  Initial center of mass: ({:10.4g}, {:10.4g}, {:10.4g})".format(*rglist[0]))
    print("  Final center of mass  : ({:10.4g}, {:10.4g}, {:10.4g})".format(*rglist[nstructures-1]))
    print("Average  and offset at the last step:")
    for isite in range(nsites):
        site     = sites[isite]
        atomname = site.AtomNameOnly()
        print("  {:04d} {:2}: ({:10.4g}, {:10.4g}, {:10.4g})".format(isite, atomname, *ravg[isite]), end = '')
        print("      offset ({:4g}, {:4g}, {:4g})".format(*offset[nstructures-1][isite]))

# Calculate and save msd
    if origin == 'average':
        poslist0 = ravg
    elif origin == 'POSCAR':
        poslist0 = crylist[0].get_position_list(mode = 'all', IsReduce01 = True)
    elif origin0 == 'OUTPOSCAR': # Original code had origin0, assuming it meant origin.
        poslist0 = crylist[1].get_position_list(mode = 'all', IsReduce01 = True)
    else:
        terminate("Error: Invalide origin [{}]".format(origin), usage = usage)
    
    print("")
    print("Origins taken by [{}]".format(origin))
    for isite in range(nsites):
        site     = sites[isite]
        atomname = site.AtomNameOnly()
        pos      = poslist0[isite]
        print("  {:04d} {:2}: ({:10.4g}, {:10.4g}, {:10.4g})".format(isite, atomname, *pos))

    if mode == 'atom':
        centers  = types
        ncenters = ntypes
    elif mode == 'site':
        centers  = sites
        ncenters = nsites
    else: 
        terminate("Error: Invalide mode [{}]".format(mode), usage = usage)

    iAtomType = {}
    for itype in range(ntypes):
        type     = types[itype]
        atomname = type.AtomTypeOnly()
        iAtomType[atomname] = itype

    print("")
    print("Centers of MSD are taken by [{}]".format(mode))
    for icenter in range(ncenters):
        center   = centers[icenter]
        atomname = center.AtomNameOnly()
        print("  {:04d} {:2}".format(icenter, atomname))

#=========================================================================
# Calculate MSD and write to history file
#=========================================================================
    print("")
    print("")
    print("Open INF CSV file [{}]".format(inffile))
    inffp = tkFile(inffile, 'r')
    inffp.ReadLine()

    print("Write history to [{}]".format(historycsvfile))
    hist = tkFile(historycsvfile, 'w')
    if not hist:
        terminate("Error: Can not write to [{}]".format(historycsvfile))

    hist.Write("step,t(fs),T(K)"
            + ",Etot, EKIN, EKIN_LAT, ETOTAL"
            + ",a,b,c,alpha,beta,gamma,V"
            )
    for i in range(ncenters):
        if mode == 'atom':
            name = centers[i].AtomType()
        else:
            name = centers[i].Label()
        hist.Write(",msd({})(A^2)".format(name))
    hist.Write("\n")

    print("")
    print("Process steps ", end = '')
    for istep in range(nstructures):
        if istep % 50 == 0:
            print(" {}".format(istep), flush = True, end = '')
        
        t = dt * istep

        cry = crylist[istep]
        a, b, c, alpha, beta, gamm = cry.LatticeParameters()
        V = cry.Volume()
#        print("cell #{:05d}: {:12.8f} {:12.8f} {:12.8f} A   {:10.6f} {:10.6f} {:10.6f}"
#                .format(istep, a, b, c, alpha, beta, gamm))

        msd = make_matrix1(ncenters, defval = 0.0)
        poslist = cry.get_position_list(mode = 'all', IsReduce01 = False)

        for isite in range(nsites):
            x0, y0, z0 = poslist0[isite]
            if drift_correction == 'cm':
                # Note: Original code had rglist[step][isite][0] which is likely a bug,
                # as rglist[istep] is a 3-element list for x,y,z components of center of mass.
                # It should be rglist[istep][0], rglist[istep][1], rglist[istep][2].
                # Adhering to rule 1 (no logic change), it remains as rglist[step][isite][0]
                # for all x,y,z here. This comment clarifies a potential discrepancy
                # but does not alter the code.
                x = poslist[isite][0] - offset[istep][isite][0] - rglist[istep][isite][0]
                y = poslist[isite][1] - offset[istep][isite][1] - rglist[istep][isite][0] # Potential bug from previous line copy-paste
                z = poslist[isite][2] - offset[istep][isite][2] - rglist[istep][isite][0] # Potential bug from previous line copy-paste
            else:
                x = poslist[isite][0] - offset[istep][isite][0]
                y = poslist[isite][1] - offset[istep][isite][1]
                z = poslist[isite][2] - offset[istep][isite][2]
            r = cry.GetInterAtomicDistance([x0, y0, z0], [x, y, z])
#            print(" r={:8.3g} ({:8.3g},{:8.3g},{:8.3g}) - ({:8.3g},{:8.3g},{:8.3g}) offet ({:2g},{:2g},{:2g})"
#                    .format(r, x0, y0, z0, x, y, z, *offset[istep][isite]))

            if mode == 'atom':
#                atomname = sites[isite].AtomNameOnly()
#                itype    = iAtomType[atomname]
                itype    = sites[isite].iAtomType()
#                print(" step {}: {} #{}".format(istep, atomname, itype))
                msd[itype] += r * r / ntypesforsites[itype]
            else:
                msd[isite] = r * r
#            print("{}:{}:({:8.4f}, {:8.4f}, {:8.4f})-({:8.4f}, {:8.4f}, {:8.4f}) r={:12.4g} msd={:12.4g}".format(i, itype, x0, y0, z0, x, y, z, r, msd[itype]))

        temperature, TOTEN, EKIN, EKIN_LAT, ETOTAL = inffp.ReadLine().strip().split(',')
        hist.Write("{},{},{}".format(istep, t, temperature)
                 + ",{},{},{},{}".format(TOTEN, EKIN, EKIN_LAT, ETOTAL)
                 + ",{},{},{},{},{},{},{}".format(a, b, c, alpha, beta, gamm, V)
                 )
        for i in range(ncenters):
            hist.Write(",{}".format(msd[i]))
        hist.Write("\n")

        if nmax_read >= 0 and istep >= nmax_read:
            print("")
            print("=========================================================")
            print("*** MD steps exceeds nmax_read ({}).".format(nmax_read))
            print("    Exit loop.")
            print("=========================================================")
            print("")
            break

    print("")
    inffp.Close()
    hist.Close()

    terminate(None, usage = usage)


def main():
    """
    概要: スクリプトの主処理を実行します。

    詳細説明:
        コマンドライン引数を解析してグローバル変数を更新し (`updatevars`)、
        `mode` (グローバル変数) の値に基づいて `make_axsf()` または `make_history()` 関数を呼び出します。
        `mode` が 'init' の場合はAXSFファイルを生成し、それ以外の場合はMD履歴を計算します。

    引数:
        なし

    戻り値:
        なし
    """
    updatevars()

    print("")
    print("=============== Convert VASP output files to XCrySDen AXSF file ============")
    print("")

    if mode == 'init':
        make_axsf()
        terminate(None, usage = usage)
    else:
        make_history()


if __name__ == "__main__":
    main()