"""
VASPのOUTCARファイルからXCrySDen AXSFアニメーションファイルを生成するスクリプト。

概要:
    VASPのOUTCARファイル（分子動力学シミュレーション結果）を解析し、XCrySDenで可視化可能な
    AXSFアニメーションファイルを生成します。各ステップの結晶構造（格子ベクトル、原子座標）
    および原子にかかる力（または速度）をAXSF形式で出力します。

詳細説明:
    このスクリプトは、VASPの計算ディレクトリ内のOUTCAR、INCAR、POSCAR、CONTCARファイルを
    読み込み、以下の処理を実行します。
    1.  POSCARから初期構造、CONTCARから最終構造を読み込み、それぞれのCIFファイルを生成します。
    2.  INCARファイルから`SYSTEM`名と`POTIM`（時間刻み）を抽出し、出力ファイル名などに利用します。
    3.  OUTCARファイルを解析し、シミュレーション中に記録されたすべての結晶構造ステップを抽出し、
        それらを順次AXSFファイルに書き込みます。
    4.  AXSFファイルには、各ステップの格子ベクトル、原子のカートesian座標、および原子にかかる力
        （OUTCARに力が記録されていない場合は速度）が出力されます。

使用方法:
    コマンドラインからVASPのOUTCARファイル（またはその親ディレクトリ）のパスを指定して実行します。
    例: python outcar2axsf.py OUTCAR
    例: python outcar2axsf.py /path/to/vasp_calculation_directory

関連リンク:
    :doc:`outcar2axsf_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
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.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'
xsffile  = None


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

    概要:
        コマンドラインでのスクリプトの実行方法をユーザーに提示します。
    """
    global CAR_path

    print("")
    print("Usage:")
    print("  (a)  python {} outcar_path".format(sys.argv[0]))
    print("     ex: python {} {}".format(sys.argv[0], CAR_path))
    
    
def updatevars():
    """
    コマンドライン引数に基づいてグローバル変数を更新します。

    概要:
        コマンドラインで指定された引数からOUTCARのパスを取得し、
        それに基づいてAXSF出力ファイル名を決定します。

    詳細説明:
        `sys.argv`から最初の引数を`CAR_path`として取得します。
        `CAR_path`からディレクトリ、ファイル名、拡張子を分離し、
        AXSFファイルの名前を`.xsf`拡張子で生成し、`xsffile`に設定します。
    """
    global CAR_path, xsffile

    CAR_path = getarg(1, CAR_path)

    dirname, basename, filebody, ext = SplitFilePath(CAR_path)
    xsffile = os.path.join(dirname, filebody + '.xsf')


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

    概要:
        新しいAXSFファイルの先頭にアニメーションのステップ数と結晶であることを示すヘッダー情報を記述します。

    詳細説明:
        指定されたファイルパスでAXSFファイルを書き込みモードで開き、
        `ANIMSTEPS nStep`と`CRYSTAL`の行を書き込みます。
        ファイルオープンに失敗した場合はNoneを返します。

    :param outfile: 出力するAXSFファイルのパス。
    :type outfile: str
    :param nStep: アニメーションの総ステップ数。
    :type nStep: int
    :returns: ファイルオープンおよび書き込みに成功した場合1、失敗した場合None。
    :rtype: int or 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):
    """
    特定のステップの結晶構造データをAXSFファイルに追記します。

    概要:
        与えられたtkCrystalオブジェクトから、そのステップの格子ベクトルと原子座標、および力を
        AXSFファイルの形式に従って追記します。

    詳細説明:
        指定されたAXSFファイルを追記モードで開き、以下の情報を書き込みます。
        -   `PRIMVEC` と `CONVVEC` のヘッダー、および3つの格子ベクトル。
        -   `PRIMCOORD` のヘッダー、全原子数、および各原子の原子番号、カートesian座標、力（または速度）。
        力が取得できない場合は、力が0として書き込まれます。

    :param outfile: 書き込み対象のAXSFファイルのパス。
    :type outfile: str
    :param iCycle: 現在のステップのインデックス（アニメーションのフレーム番号）。
    :type iCycle: int
    :param cry: 書き込む結晶構造データを持つtkCrystalオブジェクト。
    :type cry: tkCrystal
    :returns: ファイルオープンおよび書き込みに成功した場合1、失敗した場合None。
    :rtype: int or 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()
        
        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

#	my $xc = new XCrySDen;
#	$xc->WriteXSFAnimationFileHeader($XSFFile, $nStep);
#		$xc->WriteXSFFileFromCrystal($XSFFile, $iStep, $Crystal);

def outcar2axsf():
    """
    VASPのOUTCARファイルからAXSFアニメーションファイルを生成するメインロジック。

    概要:
        VASP計算ディレクトリ内のOUTCAR、INCAR、POSCAR、CONTCARファイルを読み込み、
        必要な情報を抽出し、XCrySDen形式のAXSFアニメーションファイルと初期/最終構造のCIFファイルを生成します。

    詳細説明:
        1.  VASP関連ファイルのパスを特定し、表示します。
        2.  POSCARおよびCONTCARから初期/最終結晶構造を読み込み、その格子定数を表示します。
        3.  INCARから`SYSTEM`名と`POTIM`を読み込み、表示します。
        4.  初期構造と最終構造をCIFファイルとして出力します（ファイル名は`SYSTEM`-initial.cif, `SYSTEM`-final.cif）。
        5.  OUTCARファイルから読み取れる結晶構造の総ステップ数を取得し、AXSFファイルのヘッダーを書き込みます。
        6.  OUTCARファイルを再度開き、`read_next_crystalstructure`メソッドを使用して各ステップの結晶構造を順次読み込みます。
        7.  読み込んだ各ステップの結晶構造データ（`tkCrystal`オブジェクト）を`WriteXSFFileFromCrystal`関数に渡し、
            AXSFファイルに追記していきます。
        8.  処理完了後、スクリプトを終了します。

    グローバル変数:
        :global CAR_path: VASP計算ディレクトリまたはOUTCARファイルのパス。
        :global xsffile: 出力されるAXSFファイルのパス。
    """
    global CAR_path, xsffile
    
    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)
    print("CAR dir : ", CAR_path)
    print("INCAR   : ", INCAR_path)
    print("POSCAR  : ", POSCAR_path)
    print("CONTCAR : ", CONTCAR_path)
    print("OUTCAR  : ", OUTCAR_path)
    print("AXSF file: ", xsffile)

    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))

    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("")
    print("Read [{}]".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'
    historycsvfile  = sample_name + '-history.csv'

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

    nstructures = vasp.read_n_crystalstructures(OUTCAR_path)
    print("# of crystal structures:", nstructures)

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

    print("")
    print("Read crystal structures from [{}]".format(OUTCAR_path))
    fp = tkFile(OUTCAR_path, 'r')
    if not fp:
        terminate("Error: Can not read [{}]".format(OUTCAR_path), usage = usage) # Modified outcarpath to OUTCAR_path

    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("# of atom sites for atom types = ", nsitesfortypes)

    count = 0
    while 1:
        inf, cry = vasp.read_next_crystalstructure(fp, initial_cry, IsReduce01 = True)
        if not cry:
            break

        WriteXSFFileFromCrystal(xsffile, count, cry)

        count += 1

    fp.Close()

    terminate(None, usage = usage)


def main():
    """
    スクリプトのメインエントリーポイント。

    概要:
        コマンドライン引数の処理と、OUTCARからAXSFファイルへの変換処理を実行します。

    詳細説明:
        1.  `updatevars()` を呼び出し、コマンドライン引数に基づいてグローバル変数（`CAR_path`, `xsffile`）を初期化します。
        2.  処理の開始メッセージを表示します。
        3.  `outcar2axsf()` 関数を呼び出し、VASPファイルの解析とAXSFファイルの生成を実行します。
    """
    updatevars()

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

    outcar2axsf()


if __name__ == "__main__":
    main()