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:/Programs/python/lib")
sys.path.append("d:/Programs/python/lib")

from tkfile import tkFile
from tkutils import IsDir, IsFile, SplitFilePath
from tksci.tksci import Reduce01, Round
from tkutils import terminate, pint, pfloat, getarg, getintarg, getfloatarg
from tksci.tkmatrix import make_matrix1, make_matrix2, make_matrix3
from tkcrystal.tkcif import tkCIF, tkCIFData
from tkcrystal.tkcrystal import tkCrystal
from tkcrystal.tkatomtype import tkAtomType


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

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

freduce01 = 1

#=============================
# Treat argments
#=============================
def usage():
    global cfgfile, xyzfile
    global freduce01

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

def updatevars():
    global cfgfile, xyzfile, ciffile
    global freduce01

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

    xyzfile   = getarg     (1, xyzfile)
    freduce01 = getintarg  (2, freduce01)
    header, ext = os.path.splitext(xyzfile)
    filebody    = os.path.basename(header)
    xyzfile     = filebody + '.xyz'
    cfgfile     = filebody + '.cfg'
    ciffile     = filebody + '.cif'


def read_xyzfile(cfgfile, xyzfile):
    cfg = tkFile(cfgfile, 'r')
    if not cfg:
        terminate("Error in read_xyzfile: Can not read [{}]".format(cfgfile), usage = usage)

    line        = cfg.ReadLine()
    sample_name = cfg.ReadLine().strip()
    print("sample name: [{}]".format(sample_name))

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

    cfg.Close()

    for i in range(3):
        for j in range(3):
            aij[i][j] = pfloat(aij[i][j])

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

    xyz = tkFile(xyzfile, 'r')
    if not xyz:
        terminate("Error in read_xyzfile: Can not read [{}]".format(xyzfile), usage = usage)

    nsites = pint(xyz.ReadLine())
    print("nsites = ", nsites)

# blank line
#    line = xyz.ReadLine()

    site = []
    idx = 0
    for i in range(nsites):
        idx += 1
        line = xyz.ReadLine()
        if not line:
            break

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

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

    xyz.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(len(site)):
        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)
#        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])
    
    cry.ExpandCoordinates()

#    cry.PrintInf()

    return cry


def xyz2cif():
    global cfgfile, xyzfile

    print("CFG file: {}".format(cfgfile))
    print("XYZ file: {}".format(xyzfile))

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

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

    if not cry:
        terminate("Error: Could not get crystal object from infile [{}]".format(infile), 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()

    print("")
    print("=============== Convert XYZ+CFG files to CIF file ============")
    xyz2cif()


if __name__ == "__main__":
    main()