"""
概要:
    CFGファイルをCIFファイルに変換するスクリプトです。
詳細説明:
    Open-CFG形式のファイルから結晶構造情報を読み込み、CIF (Crystallographic Information File) 形式に変換して出力します。
    このスクリプトは、格子ベクトルの読み込み、原子サイトの座標処理
    （必要に応じて[0,1)範囲への丸め込み、指定範囲外の除外）をサポートします。
関連リンク:
    cfg2cif_usage
"""
import os
import sys
import shutil
import glob
import csv
import numpy as np
from numpy import exp, log, sin, cos, tan, arcsin, arccos, arctan, pi
from scipy.interpolate import interp1d
from pprint import pprint
from matplotlib import pyplot as plt

sys.path.append("c:/git/tkProg/tklib/python")
sys.path.append("d:/git/tkProg/tklib/python")

from tklib.tkfile import tkFile
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.tkatomtype import tkAtomType


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

atomtypes = ['O', 'Zn', 'Ga', 'In']
xlim, ylim, zlim = 1.0, 1.0, 1.0
#xlim, ylim, zlim = 0.7, 0.7, 0.7

cfgfile = 'STD0.cfg'
ciffile = 'STD0.cif'

freduce01 = 1


#=============================
# Treat argments
#=============================
def usage():
    """
    概要:
        スクリプトのコマンドライン引数の利用方法を表示します。
    詳細説明:
        スクリプト実行時に指定可能な引数（入力CFGファイル名、座標の[0,1)範囲への削減フラグ、
        座標の出力範囲リミット）とその意味を示します。
    """
    global cfgfile
    global freduce01, xlim, ylim, zlim

    print("")
    print("Usage:")
    print("  (i) python {} cfgfile (reduce01 xilm ylim zlim)".format(sys.argv[0]))
    print("      reduce01: Reduce fractional coordinates to [0, 1} if reduce01 == 1")
    print("      ex: python {} {} {} {} {} {}"
                    .format(sys.argv[0], cfgfile, freduce01, xlim, ylim, zlim))

def updatevars():
    """
    概要:
        コマンドライン引数を解析し、グローバル設定変数を更新します。
    詳細説明:
        sys.argv から引数を取得し、cfgfile, freduce01, xlim, ylim, zlim などの
        グローバル変数を設定します。入力CFGファイル名に基づいて出力CIFファイル名も決定します。
        引数の取得には getarg, getintarg, getfloatarg を使用します。
    """
    global cfgfile, ciffile
    global freduce01, xlim, ylim, zlim

    argv = sys.argv
#    if len(argv) == 1:
#        terminate(usage = usage)

    cfgfile   = getarg     (1, cfgfile)
    freduce01 = getintarg  (2, freduce01)
    xlim      = getfloatarg(3, xlim)
    ylim      = getfloatarg(4, ylim)
    zlim      = getfloatarg(5, zlim)

    header, ext = os.path.splitext(cfgfile)
    filebody    = os.path.basename(header)
    cfgfile     = filebody + '.cfg'
    ciffile     = filebody + '.cif'


def read_cfgfile(cfgfile):
    """
    概要:
        指定されたCFGファイルを読み込み、tkCrystalオブジェクトを構築します。
    詳細説明:
        CFGファイルからサンプル名、総原子数、格子ベクトル、各原子サイトの座標と原子種を抽出します。
        格子ベクトルの[0]行目が2倍されている点に注意が必要です。
        抽出された情報をもとにtkCrystalオブジェクトを初期化し、原子サイトを追加します。
        freduce01が真の場合、原子座標は[0,1)範囲に正規化され、
        xlim, ylim, zlimで指定された範囲外の原子は除外されます。
    引数:
        :param cfgfile: 読み込むCFGファイルのパス。
        :type cfgfile: str
    戻り値:
        :returns: 構築されたtkCrystalオブジェクト。ファイルの読み込みに失敗した場合はNoneを返す。
        :rtype: tkCrystal or None
    """
    cfg = tkFile(cfgfile, 'r')
    if not cfg:
        terminate("Error in read_cfgfile: Can not read [{}]".format(cfgfile), usage = usage)

    line        = cfg.ReadLine()
    sample_name = cfg.ReadLine().strip()
    print("sample name: [{}]".format(sample_name))
    
    line = cfg.SkipTo(r"molecules\s+of\s+all")
    ret = cfg.Search(r"(\d+)", line)
    natoms = pint(ret[1])

    line = cfg.SkipTo("Defining vectors")
    aij = np.empty([3, 3])
    aij[0] = cfg.ReadLine().split()
#    aij[1] = cfg.ReadLine().split()
#    aij[2] = cfg.ReadLine().split()
    for i in range(3):
        aij[0][i] = pfloat(aij[0][i]) * 2.0
    aij[1] = [0.0, aij[0][0], 0.0]
    aij[2] = [0.0, 0.0, aij[0][0]]

    print("")
    print("Lattice vectors:")
    for i in range(3):
        print("  {:16.12f}  {:16.12f}  {:16.12f}".format(*aij[i]))

    ntypes = []
    ntot = 0
    while 1:
        line = cfg.SkipTo(r"molecules\s+of\s+type")
        if not line:
            break

        ret = cfg.Search(r"(\d+)", line)
        n = pint(ret[1])
        ntypes.append(n)
        
        ntot += n
        if ntot >= natoms:
            break

    print("")
    print("natoms=", natoms)
    print("  ntypes=", ntypes) 

# skip two lines after the "molecules of type" line
    cfg.ReadLine()
    cfg.ReadLine()

    site = []
    idx = 0
    idxtype = 0
    itype = 0
    while 1:
        line = cfg.ReadLine()
        if not line:
            break

        ret = line.split()
# skip blank line
        if ret is None or len(ret) < 3:
            continue

        name = atomtypes[itype]
        xc, yc, zc = ret
        xc = pfloat(xc)
        yc = pfloat(yc)
        zc = pfloat(zc)
        print("{:04d} {:2} #{:04d} ({:8.4f}, {:8.4f}, {:8.4f})".format(idx+1, name, idxtype+1, xc, yc, zc))
        site.append([name, xc, yc, zc])

        idx += 1
        idxtype += 1
        if idxtype >= ntypes[itype]:
            itype += 1
            idxtype = 0

    cfg.Close()

    print("")
    print("Build Crystal object:")
    cry = tkCrystal()
    cry.SetSampleName(sample_name)
    cry.SetCrystalName(sample_name)
    cry.SetLatticeVectors(aij)
    latt = cry.LatticeParameters()
    for i in range(natoms):
        if i % 100 == 0:
            print("{} ".format(i))
        x, y, z = site[i][1], site[i][2], site[i][3]
#        x, y, z = cry.CartesianToFractional(site[i][1], site[i][2], site[i][3])
        if freduce01:
            x, y, z = Reduce01(x), Reduce01(y), Reduce01(z)
        if x > xlim or y > ylim or z > zlim:
            continue
#        print("xyz {}:{}: {},{},{} => {},{},{}".format(i+1, site[i][0], site[i][1], site[i][2], site[i][3], x, y, z))
        cry.AddAtomSite(name = site[i][0], pos = [x, y, z])

    print("")
    
    cry.ExpandCoordinates()

#    cry.PrintInf()

    return cry


def cfg2cif():
    """
    概要:
        グローバル変数で指定されたCFGファイルを読み込み、CIFファイルに変換して保存します。
    詳細説明:
        read_cfgfile関数を呼び出してCFGファイルから結晶情報を取得し、tkCrystalオブジェクトを生成します。
        その後、そのtkCrystalオブジェクトを使用してtkCIFDataオブジェクトを生成し、
        最終的にCIFファイルとして情報をciffileに書き出します。
        処理の最後に terminate 関数が呼ばれ、スクリプトの利用方法が表示されます。
    """
    global cfgfile
    
    print("CFG file: {}".format(cfgfile))

    print("")
    print("Read [{}]".format(cfgfile))

#    if not tkutils.IsFile(infile):
#        terminate("Error: Invalid infile [{}]".format(infile), usage = usage)
    cry = read_cfgfile(cfgfile)

    if not cry:
        terminate("Error: Could not get crystal object from infile [{}]".format(cfgfile), usage = usage)

    print("")
    print("==============================================")
    print("    Crystal inf")
    print("==============================================")
    cry.PrintInf()

    print("")
    print("Save to [{}]".format(ciffile))
    cif = tkCIFData()
    cif.CreateCIFFileFromCCrystal(cry, ciffile)

    terminate(usage = usage)


def main():
    """
    概要:
        スクリプトのエントリポイントです。
    詳細説明:
        コマンドライン引数の処理 (updatevars) と、
        CFGからCIFへの変換処理 (cfg2cif) を順に実行します。
        スクリプトの開始時と終了時にメッセージを出力します。
    """
    updatevars()

    print("")
    print("=============== Convert CFG files to CIF file ============")
    cfg2cif()


if __name__ == "__main__":
    main()